// t5de-renderer.exe: the renderer process. Started by the adapter (t5deRenderEngine.dll), one for the whole client:
// receives each 3D view's drawing as commands through a shared memory channel of its own, and draws it with Diligent
// into a child window it places over the view in the client's window.

#include <windows.h>

#include <d3d11.h>
#include <dxgi.h>
#include <shellapi.h>

#include <algorithm>
#include <atomic>
#include <cstdint>
#include <cstring>
#include <filesystem>
#include <format>
#include <memory>
#include <mutex>
#include <optional>
#include <string>
#include <string_view>
#include <thread>
#include <vector>

#include "EngineFactoryD3D11.h"
#include "RefCntAutoPtr.hpp"
#include "RenderDeviceD3D11.h"
#include "log.hpp"
#include "scene.hpp"
#include "t5de/protocol/channel.hpp"

namespace {

using namespace Diligent;
using t5de::protocol::Channel;
using t5de::protocol::View;

using t5de::log::info;

// Diligent's messages (shader compilation errors, invalid states, ...), into the log.
void DILIGENT_CALL_TYPE diligent_message(DEBUG_MESSAGE_SEVERITY severity, const Char* message, const Char* function,
                                         const Char*, int)
{
    static constexpr const char* severities[] = {"info", "warning", "error", "fatal error"};
    info("Diligent {}: {}{}{}", severities[severity < 4 ? severity : 3], function != nullptr ? function : "",
         function != nullptr ? ": " : "", message != nullptr ? message : "");
}

struct Arguments {
    std::wstring control;
    DWORD parent = 0;
};

Arguments parse_arguments()
{
    Arguments arguments;
    int count = 0;
    wchar_t** argv = CommandLineToArgvW(GetCommandLineW(), &count);
    for (int i = 1; i + 1 < count; i += 2) {
        const std::wstring_view name = argv[i];
        if (name == L"--control") {
            arguments.control = argv[i + 1];
        } else if (name == L"--parent") {
            arguments.parent = static_cast<DWORD>(std::wcstoul(argv[i + 1], nullptr, 10));
        }
    }
    LocalFree(argv);
    return arguments;
}

std::string narrow(const std::wstring& text) { return std::filesystem::path(text).string(); }

LRESULT CALLBACK window_procedure(HWND window, UINT message, WPARAM wparam, LPARAM lparam)
{
    if (message == WM_ERASEBKGND) {
        return 1; // we draw all of it
    }
    return DefWindowProcW(window, message, wparam, lparam);
}

// A child of the client's window. Disabled, so mouse input goes to the client's window under it.
HWND create_child_window(HWND parent)
{
    static const ATOM windowClass = [] {
        const WNDCLASSEXW description{.cbSize = sizeof(WNDCLASSEXW),
                                      .lpfnWndProc = window_procedure,
                                      .hInstance = GetModuleHandleW(nullptr),
                                      .hCursor = LoadCursorW(nullptr, IDC_ARROW),
                                      .lpszClassName = L"T5DERenderer"};
        return RegisterClassExW(&description);
    }();
    return CreateWindowExW(0, MAKEINTATOM(windowClass), L"", WS_CHILD | WS_VISIBLE | WS_DISABLED | WS_CLIPSIBLINGS, 0,
                           0, 1, 1, parent, nullptr, GetModuleHandleW(nullptr), nullptr);
}

// The view, as the adapter set it, or as %TEMP%\t5de-renderer-view.txt says while the client runs (off, split, full
// or compare): for comparing the renderers without restarting the client. Read every second or so.
View current_view(View adapter)
{
    static std::optional<View> view;
    static ULONGLONG checked = 0;
    if (GetTickCount64() - checked >= 1000) {
        checked = GetTickCount64();
        view.reset();
        wchar_t temp[MAX_PATH + 1];
        const DWORD length = GetTempPathW(MAX_PATH + 1, temp);
        const HANDLE file = length == 0 || length > MAX_PATH
                                ? INVALID_HANDLE_VALUE
                                : CreateFileW((std::filesystem::path(temp) / L"t5de-renderer-view.txt").c_str(),
                                              GENERIC_READ, FILE_SHARE_READ | FILE_SHARE_WRITE | FILE_SHARE_DELETE,
                                              nullptr, OPEN_EXISTING, FILE_ATTRIBUTE_NORMAL, nullptr);
        if (file != INVALID_HANDLE_VALUE) {
            char text[16] = {};
            DWORD read = 0;
            ReadFile(file, text, sizeof(text) - 1, &read, nullptr);
            CloseHandle(file);
            const std::string_view setting(text, std::strcspn(text, "\r\n "));
            if (setting == "off") {
                view = View::Off;
            } else if (setting == "split") {
                view = View::Split;
            } else if (setting == "full") {
                view = View::Full;
            } else if (setting == "compare") {
                view = View::Compare;
            }
        }
    }
    return view.value_or(adapter);
}

// Shader changes to try while the client runs: %TEMP%\t5de-renderer-shaders\fixed_function.hlsl, if there is
// one, replaces the built-in shaders, and is read again when it changes (checked every second or so). Returns the
// new source when it has changed.
std::optional<std::string> changed_shaders()
{
    static ULONGLONG checked = 0;
    static std::filesystem::file_time_type loaded{};
    if (GetTickCount64() - checked < 1000) {
        return std::nullopt;
    }
    checked = GetTickCount64();
    std::error_code error;
    const auto path = std::filesystem::temp_directory_path(error) / L"t5de-renderer-shaders" / L"fixed_function.hlsl";
    const auto modified = std::filesystem::last_write_time(path, error);
    if (error || modified == loaded) {
        return std::nullopt;
    }
    const HANDLE file = CreateFileW(path.c_str(), GENERIC_READ, FILE_SHARE_READ | FILE_SHARE_WRITE, nullptr,
                                    OPEN_EXISTING, FILE_ATTRIBUTE_NORMAL, nullptr);
    if (file == INVALID_HANDLE_VALUE) {
        return std::nullopt;
    }
    std::string source(static_cast<std::size_t>(std::filesystem::file_size(path, error)), ' ');
    DWORD read = 0;
    ReadFile(file, source.data(), static_cast<DWORD>(source.size()), &read, nullptr);
    CloseHandle(file);
    source.resize(read);
    loaded = modified;
    info("shaders reloaded from {}", path.string());
    return source;
}

// Where the child window goes in the client's 3D view, and where in the 3D view its scene starts (offset).
struct Placement {
    RECT window{};
    int offset = 0;
};

Placement placement(HWND parent, View view)
{
    Placement place;
    GetClientRect(parent, &place.window);
    if (view == View::Split || view == View::Compare) {
        place.window.left = place.window.right / 2;
    }
    place.offset = view == View::Compare ? 0 : place.window.left;
    if (view == View::Off) {
        place.window = {};
    }
    return place;
}

// One of the client's 3D views: an engine's channel, and the child window its frames are drawn into. A thread of its
// own takes the channel's chunks as soon as they come (so the client never waits for drawing) and queues them; the main
// thread draws the newest complete frame, carrying out older ones' changes without drawing them.
class ViewWindow {
public:
    // Opens the channel named name and sets up its window; null if the channel is gone or it can't.
    static std::unique_ptr<ViewWindow> open(const std::wstring& name, IEngineFactoryD3D11* factory,
                                            IRenderDevice* device, IDeviceContext* context,
                                            const std::optional<std::string>& shaders)
    {
        std::unique_ptr<ViewWindow> view(new ViewWindow());
        view->name_ = narrow(name.substr(name.rfind(L'-') + 1));
        view->mapping_ = OpenFileMappingW(FILE_MAP_ALL_ACCESS, FALSE, name.c_str());
        if (view->mapping_ == nullptr) {
            return nullptr; // its engine is gone already
        }
        view->channel_ = static_cast<Channel*>(MapViewOfFile(view->mapping_, FILE_MAP_ALL_ACCESS, 0, 0, 0));
        view->ready_ = OpenEventW(SYNCHRONIZE, FALSE, (name + t5de::protocol::readyEventSuffix).c_str());
        view->consumed_ = OpenEventW(EVENT_MODIFY_STATE, FALSE, (name + t5de::protocol::consumedEventSuffix).c_str());
        view->read_ = OpenEventW(EVENT_MODIFY_STATE, FALSE, (name + t5de::protocol::readEventSuffix).c_str());
        if (view->channel_ == nullptr || view->ready_ == nullptr || view->consumed_ == nullptr
            || view->read_ == nullptr || view->channel_->version != t5de::protocol::version) {
            info("{}: can't open the channel", view->name_);
            return nullptr;
        }
        view->data_ = reinterpret_cast<std::byte*>(view->channel_ + 1);
        view->clientWindow_ = reinterpret_cast<HWND>(static_cast<std::uintptr_t>(view->channel_->window));
        view->window_ = create_child_window(view->clientWindow_);
        if (view->window_ == nullptr) {
            info("{}: can't create a window in {:#x}", view->name_, view->channel_->window);
            return nullptr;
        }
        SwapChainDesc swapChainDesc;
        swapChainDesc.ColorBufferFormat = TEX_FORMAT_RGBA8_UNORM; // no sRGB conversion, like the client's renderer
        swapChainDesc.DepthBufferFormat = TEX_FORMAT_D32_FLOAT;
        factory->CreateSwapChainD3D11(device, context, swapChainDesc, FullScreenModeDesc{},
                                      Win32NativeWindow{view->window_}, &view->swapChain_);
        if (!view->swapChain_) {
            info("{}: can't create a swap chain", view->name_);
            return nullptr;
        }
        view->scene_ = std::make_unique<t5de::Scene>(device, context, view->swapChain_);
        if (shaders) {
            view->scene_->setShaderSource(*shaders);
        }
        // Frames read back: the adapter waits for them, sending nothing meanwhile, so the data area is free.
        ViewWindow* self = view.get();
        view->scene_->setReadBack([self](std::span<const std::byte> pixels) {
            const std::size_t size = pixels.size() <= self->channel_->capacity ? pixels.size() : 0;
            std::memcpy(self->data_, pixels.data(), size);
            self->channel_->readSize = static_cast<std::uint32_t>(size);
            SetEvent(self->read_);
        });
        QueryPerformanceFrequency(&view->frequency_);
        view->framesReady_ = CreateEventW(nullptr, FALSE, FALSE, nullptr);
        if (view->framesReady_ == nullptr) {
            return nullptr;
        }
        view->receiver_ = std::thread([self] { self->receive(); });
        info("{}: drawing window {:#x}, view {}", view->name_, view->channel_->window,
             static_cast<int>(view->channel_->view));
        return view;
    }

    ViewWindow(const ViewWindow&) = delete;
    ViewWindow& operator=(const ViewWindow&) = delete;

    ~ViewWindow()
    {
        stopping_ = true;
        if (receiver_.joinable()) {
            receiver_.join();
        }
        if (channel_ != nullptr) {
            channel_->covering = 0;
            UnmapViewOfFile(channel_);
        }
        scene_.reset();
        swapChain_.Release();
        if (window_ != nullptr) {
            DestroyWindow(window_);
            info("{}: closed", name_);
        }
        for (HANDLE handle : {ready_, consumed_, read_, mapping_, framesReady_}) {
            if (handle != nullptr) {
                CloseHandle(handle);
            }
        }
    }

    // Set when complete frames are queued.
    [[nodiscard]] HANDLE framesReady() const { return framesReady_; }
    // Its engine is gone.
    [[nodiscard]] bool closing() const { return channel_->closing != 0; }

    void setShaderSource(const std::string& source) { scene_->setShaderSource(source); }

    // After framesReady: draws the newest complete frame queued. If the client has sent more meanwhile (drawing fell
    // behind), the older ones' changes are carried out without drawing them.
    void onFramesReady()
    {
        std::vector<std::size_t> ends;
        {
            const std::scoped_lock lock(mutex_);
            if (frameEnds_.empty()) {
                return;
            }
            const std::size_t complete = frameEnds_.back();
            frame_.assign(received_.begin(), received_.begin() + static_cast<std::ptrdiff_t>(complete));
            received_.erase(received_.begin(), received_.begin() + static_cast<std::ptrdiff_t>(complete));
            ends.swap(frameEnds_);
        }
        std::size_t begin = 0;
        for (std::size_t i = 0; i + 1 < ends.size(); ++i) {
            scene_->execute(std::span(frame_).subspan(begin, ends[i] - begin), 0, false);
            begin = ends[i];
            ++skipped_;
        }
        const std::size_t complete = ends.back();

        // Follow the client's 3D view's size.
        const View view = current_view(channel_->view);
        const auto [place, offset] = placement(clientWindow_, view);
        const int width = place.right - place.left;
        const int height = place.bottom - place.top;
        if (std::memcmp(&place, &placed_, sizeof(RECT)) != 0) {
            placed_ = place;
            if (width > 0 && height > 0) {
                SetWindowPos(window_, nullptr, place.left, place.top, width, height,
                             SWP_NOZORDER | SWP_NOACTIVATE | SWP_SHOWWINDOW);
                swapChain_->Resize(static_cast<Uint32>(width), static_cast<Uint32>(height));
                info("{}: drawing {}x{} at {} in the client's 3D view, its scene from {}", name_, width, height,
                     place.left, offset);
            } else {
                ShowWindow(window_, SW_HIDE);
                channel_->covering = 0;
                info("{}: hidden", name_);
            }
        }
        LARGE_INTEGER start, end;
        QueryPerformanceCounter(&start);
        const bool covers = view == View::Full && width > 0 && height > 0;
        if (scene_->execute(std::span(frame_).subspan(begin, complete - begin), offset, width > 0 && height > 0)) {
            swapChain_->Present(0);
            channel_->covering = covers ? 1 : 0; // the client's renderer can stop drawing (or must start again)
            QueryPerformanceCounter(&end);
            const auto seconds = [this](LONGLONG ticks) {
                return static_cast<double>(ticks) / static_cast<double>(frequency_.QuadPart);
            };
            drawTime_ += seconds(end.QuadPart - start.QuadPart);
            latency_ += seconds(end.QuadPart - static_cast<LONGLONG>(scene_->lastSentAt()));
            if (++drawn_ == 300) {
                info("{}: last 300 frames: {:.1f} ms each to draw and present, {:.1f} ms from the client's end of frame "
                     "to presented, {} skipped to catch up",
                     name_, drawTime_ * 1000 / drawn_, latency_ * 1000 / drawn_, skipped_);
                drawTime_ = 0;
                latency_ = 0;
                drawn_ = 0;
                skipped_ = 0;
            }
        }
    }

private:
    ViewWindow() = default;

    // The receiver thread: takes each chunk as it comes (after "ready"), lets the adapter send the next, and queues it;
    // sets framesReady when a frame is complete. Ends when the view closes.
    void receive()
    {
        while (!stopping_ && !closing()) {
            if (WaitForSingleObject(ready_, 100) != WAIT_OBJECT_0 || closing()) {
                continue;
            }
            const std::uint32_t size = channel_->chunkSize <= channel_->capacity ? channel_->chunkSize : 0;
            const bool endsFrame = (channel_->chunkFlags & t5de::protocol::chunkEndsFrame) != 0;
            {
                const std::scoped_lock lock(mutex_);
                received_.insert(received_.end(), data_, data_ + size);
                if (endsFrame) {
                    frameEnds_.push_back(received_.size());
                }
            }
            SetEvent(consumed_);
            if (endsFrame) {
                SetEvent(framesReady_);
            }
        }
        SetEvent(framesReady_); // wakes the main thread to close the view
    }

    std::string name_;
    HANDLE mapping_ = nullptr;
    HANDLE ready_ = nullptr;
    HANDLE consumed_ = nullptr;
    HANDLE read_ = nullptr;
    Channel* channel_ = nullptr;
    std::byte* data_ = nullptr;
    HWND clientWindow_ = nullptr;
    HWND window_ = nullptr;
    RefCntAutoPtr<ISwapChain> swapChain_;
    std::unique_ptr<t5de::Scene> scene_;
    std::vector<std::byte> frame_; // the frames being drawn
    RECT placed_{};

    std::thread receiver_;
    std::atomic<bool> stopping_ = false;
    HANDLE framesReady_ = nullptr;
    std::mutex mutex_;                  // guards received_ and frameEnds_
    std::vector<std::byte> received_;   // chunks taken, not yet drawn
    std::vector<std::size_t> frameEnds_; // where complete frames end in received_

    LARGE_INTEGER frequency_{};
    double drawTime_ = 0;
    double latency_ = 0;
    int drawn_ = 0;
    int skipped_ = 0;
};

} // namespace

int WINAPI wWinMain(HINSTANCE, HINSTANCE, PWSTR, int)
{
    const Arguments arguments = parse_arguments();
    if (arguments.control.empty()) {
        return 2;
    }
    HANDLE mapping = OpenFileMappingW(FILE_MAP_READ, FALSE, arguments.control.c_str());
    const auto* control = mapping == nullptr ? nullptr
                                             : static_cast<const t5de::protocol::Control*>(
                                                   MapViewOfFile(mapping, FILE_MAP_READ, 0, 0, 0));
    HANDLE attach = OpenEventW(SYNCHRONIZE, FALSE, (arguments.control + t5de::protocol::attachEventSuffix).c_str());
    HANDLE parent = OpenProcess(SYNCHRONIZE, FALSE, arguments.parent);
    if (control == nullptr || attach == nullptr || parent == nullptr || control->version != t5de::protocol::version) {
        return 3;
    }

    info("started for client {}", arguments.parent);
    IEngineFactoryD3D11* factory = GetEngineFactoryD3D11();
    factory->SetMessageCallback(diligent_message);
    RefCntAutoPtr<IRenderDevice> device;
    RefCntAutoPtr<IDeviceContext> context;
    EngineD3D11CreateInfo engineInfo;
    factory->CreateDeviceAndContextsD3D11(engineInfo, &device, &context);
    if (!device) {
        info("can't create a Direct3D 11 device");
        return 5;
    }
    // Show each frame as soon as it's drawn: Direct3D 11 otherwise lets 3 frames queue, which the view would trail by.
    if (RefCntAutoPtr<IRenderDeviceD3D11> d3d11{device, IID_RenderDeviceD3D11}) {
        IDXGIDevice1* dxgi = nullptr;
        if (SUCCEEDED(d3d11->GetD3D11Device()->QueryInterface(__uuidof(IDXGIDevice1), reinterpret_cast<void**>(&dxgi)))) {
            dxgi->SetMaximumFrameLatency(1);
            dxgi->Release();
        }
    }

    std::vector<std::unique_ptr<ViewWindow>> views;
    std::uint32_t attached = 0;
    std::optional<std::string> shaders; // the override, if any
    std::size_t first = 0;              // which view's event is waited on first, taking turns
    for (;;) {
        // New channels: opened in order; those whose engines are gone already are skipped.
        while (attached < control->channels) {
            if (auto view = ViewWindow::open(std::format(L"{}-{}", arguments.control, attached), factory, device,
                                             context, shaders)) {
                views.push_back(std::move(view));
            }
            ++attached;
        }
        std::erase_if(views, [](const auto& view) { return view->closing(); });
        if (auto source = changed_shaders()) {
            shaders = std::move(source);
            for (const auto& view : views) {
                view->setShaderSource(*shaders);
            }
        }

        std::vector<HANDLE> handles{parent, attach};
        const std::size_t count = std::min<std::size_t>(views.size(), MAXIMUM_WAIT_OBJECTS - 3);
        for (std::size_t i = 0; i < count; ++i) {
            handles.push_back(views[(first + i) % views.size()]->framesReady());
        }
        const DWORD result = MsgWaitForMultipleObjects(static_cast<DWORD>(handles.size()), handles.data(), FALSE, 1000,
                                                       QS_ALLINPUT);
        MSG message;
        while (PeekMessageW(&message, nullptr, 0, 0, PM_REMOVE)) {
            DispatchMessageW(&message);
        }
        if (result == WAIT_OBJECT_0) {
            break; // the client exited
        }
        if (result >= WAIT_OBJECT_0 + 2 && result < WAIT_OBJECT_0 + handles.size()) {
            const std::size_t index = (first + (result - WAIT_OBJECT_0 - 2)) % views.size();
            views[index]->onFramesReady();
            first = (index + 1) % views.size();
        }
    }

    views.clear();
    info("exiting");
    return 0;
}