#include <nng/nng.h>
#include <imgui.h>
#include <backends/imgui_impl_sdl.h>
#include <backends/imgui_impl_sdlrenderer.h>
#include <SDL2/SDL.h>

#include <algorithm>
#include <array>
#include <cerrno>
#include <cstdlib>
#include <thread>
#include <type_traits>
#include <vector>

#include <signal.h>
#include <spawn.h>
#include <sys/wait.h>
#include <unistd.h>

#include "dp_common.h"
#include "dp_util.hpp"

void dp_sdl_fatal(const char *const msg)
{
    log_fatal("%s: %s", msg, SDL_GetError());
    abort();
}

struct DoomState
{
    doomid_t id = 0;
    DP_DoomState state = DP_DS_Unknown;
    SDL_Texture *texture = nullptr;
};

static_assert(std::is_trivially_copyable<DoomState>::value,
    "DoomState must be a trivially copyable type");

struct ControllerContext
{
    nng_socket pub;
    nng_socket sub;
    SDL_Window *window;
    SDL_Renderer *renderer;
    std::vector<DoomState> dooms;
    bool quit = false;
    ExampleAppLog appLog;
    int columns = 4;
    std::array<u8, DoomScreenWidth * DoomScreenHeight * DoomBytesPerPixel> pixelBuffer;
};

struct ControllerActions
{
    int doomsToSpawn = 0;
    bool endAllDooms = false;
};

#define DOOM_EXECUTABLE "dp_doom"

void spawn_doom_posix_spawn(ControllerContext &ctx)
{
    DoomState ds;

    const char *const argv[] = { DOOM_EXECUTABLE, nullptr };

    // TODO: Close stdin and stdout? Leave them open for now to see the logging
    // output.
    // FIXME: SDL is not able to init its sound system in the doomchild. Try to ask
    // it if it's already initialized and skip if true.
    if (auto err = posix_spawn(&ds.id, DOOM_EXECUTABLE, nullptr, nullptr,
        const_cast<char *const *>(argv), nullptr))
    {
        log_error("Could not spawn doom: %s", strerror(err));
        return;
    }

    ds.texture = SDL_CreateTexture(ctx.renderer, SDL_PIXELFORMAT_RGB888,
        SDL_TEXTUREACCESS_STREAMING, DoomScreenWidth, DoomScreenHeight);

    if (!ds.texture)
        dp_sdl_fatal("SDL_CreateTexture");

    log_info("Spawned new doom, pid=%d", ds.id);
    ctx.dooms.emplace_back(ds);
}

// nng does not like this at all and panics (for good reasons)
void spawn_doom_fork(ControllerContext &ctx)
{
    auto pid = fork();

    if (pid < 0)
    {
        log_error("Could not spawn doom: %s", strerror(errno));
        return;
    }

    if (pid == 0)       // doomchild
    {
        log_info("doomchild %d cleaning the room", getpid());
        ImGui::DestroyContext();
        log_info("doomchild %d cleanup done", getpid());

        const char *const argv[] = { DOOM_EXECUTABLE, nullptr };
        execvp(DOOM_EXECUTABLE, const_cast<char *const *>(argv));
    }
    else if (pid > 0)   // doomparent
    {
        DoomState ds;
        ds.id = pid;
        ds.texture = SDL_CreateTexture(ctx.renderer, SDL_PIXELFORMAT_RGB888,
            SDL_TEXTUREACCESS_STREAMING, DoomScreenWidth, DoomScreenHeight);
        if (!ds.texture)
            dp_sdl_fatal("SDL_CreateTexture");

        log_info("Spawned new doom, pid=%d", ds.id);
        ctx.dooms.emplace_back(ds);
    }
}

inline void spawn_doom(ControllerContext &ctx)
{
    spawn_doom_posix_spawn(ctx);
}

void end_all_dooms(ControllerContext &ctx)
{
    nng_msg *msg = nullptr;
    int res = 0;

    if ((res = nng_msg_alloc(&msg, sizeof(MsgMcstCommand))))
        dp_nng_fatal("ctrl/nng_msg_alloc", res);

    auto dpmsg = DP_NNG_BODY_AS(msg, MsgMcstCommand);
    dpmsg->head.msgType = DP_MT_McstCommand;
    dpmsg->cmd = DP_DC_Endoom;

    if ((res = nng_sendmsg(ctx.pub, msg, 0)))
        dp_nng_fatal("ctrl/sendmsg", res);
}

void signal_all_dooms(ControllerContext &ctx, int signum)
{
    std::for_each(std::begin(ctx.dooms), std::end(ctx.dooms),
        [signum] (const auto &ds)
        {
            kill(ds.id, signum);
        });
}

void perform_actions(ControllerContext &ctx, const ControllerActions &actions)
{
    if (actions.doomsToSpawn)
    {
        log_info("Spawning %d new dooms", actions.doomsToSpawn);
        for (int i=0; i<actions.doomsToSpawn; ++i)
            spawn_doom(ctx);
    }

    if (actions.endAllDooms)
    {
        log_info("Telling all dooms to quit");
        end_all_dooms(ctx);
    }
}

void check_on_dooms(ControllerContext &ctx)
{
    pid_t pid = 0;

    // This works for dooms forked from the controller, not for externally
    // started dooms.
    do
    {
        int wstatus = 0;

        if (pid = waitpid(0, &wstatus, WNOHANG); pid > 0)
        {
            auto ds = find_in_container(ctx.dooms, [pid] (const auto &ds) { return ds.id == pid; });

            assert(ds != std::end(ctx.dooms));

            if (ds != std::end(ctx.dooms))
            {
                if (WIFEXITED(wstatus))
                    log_info("doom(%d) exited with status %d", pid, WEXITSTATUS(wstatus));
                else if (WIFSIGNALED(wstatus))
                    log_warn("doom#(%d) got killed by signal %d", pid, WTERMSIG(wstatus));

                SDL_DestroyTexture(ds->texture); // TODO: use the destructor to do this
                ctx.dooms.erase(ds);
            }
        }
    } while (pid > 0);
}

void do_networking(ControllerContext &ctx)
{
    // Set to true if we receive at least on DP_DS_Ready DoomState update. Then
    // a single DP_DC_RunDoom command is broadcast.
    bool sendRunDoom = false;

    // Limit the max time we spend doing network stuff.
    static const auto MaxNetworkingTime = std::chrono::milliseconds(10);
    auto tStart = std::chrono::steady_clock::now();

    while (true)
    {
        if (auto elapsed = std::chrono::steady_clock::now() - tStart;
            elapsed >= MaxNetworkingTime)
        {
            break;
        }

        nng_msg *msg = nullptr;

        if (auto res = dp_recv_new_msg_nonblock(ctx.sub, &msg))
        {
            if (!dp_nng_is_timeout(res))
                dp_nng_fatal("ctrl/recvmsg", res);
            break; // timeout
        }

        auto msgBase = DP_NNG_BODY_AS(msg, MessageBase);

        if (msgBase->msgType == DP_MT_DoomState)
        {
            auto msgDoomState = DP_NNG_BODY_AS(msg, MsgDoomState);

            // Check if we know this doom. If it was externally started register
            // it in ctx.dooms.
            auto pid = msgDoomState->doomId;
            auto dit = find_in_container(ctx.dooms, [pid] (const auto &ds) { return ds.id == pid; });

            if (dit != std::end(ctx.dooms))
            {
                dit->state = msgDoomState->doomState;
            }
            else
            {
                DoomState ds;
                ds.id = pid;
                ds.state = msgDoomState->doomState;
                ds.texture = SDL_CreateTexture(ctx.renderer, SDL_PIXELFORMAT_RGB888,
                                               SDL_TEXTUREACCESS_STREAMING, DoomScreenWidth, DoomScreenHeight);

                if (!ds.texture)
                    dp_sdl_fatal("SDL_CreateTexture");

                log_info("Registered external doom, pid=%d", ds.id);
                ctx.dooms.emplace_back(ds);
            }

            if (msgDoomState->doomState == DP_DS_Ready)
                sendRunDoom = true;
        }
        else if (msgBase->msgType == DP_MT_DoomFrame)
        {
            auto msgDoomFrame = DP_NNG_BODY_AS(msg, MsgDoomFrame);
            auto pid = msgDoomFrame->doomId;
            auto dit = find_in_container(ctx.dooms, [pid] (const auto &ds) { return ds.id == pid; });

            if (dit != std::end(ctx.dooms))
            {
                auto &ds = *dit;
                SDL_UpdateTexture(ds.texture, nullptr, msgDoomFrame->frame, DoomFramePitch);
            }
            else
                log_warn("Received DoomFrame from unregistered doom, pid=%d", pid);
        }

        nng_msg_free(msg);
    }

    if (sendRunDoom)
    {
        nng_msg *msg = nullptr;
        int res = 0;

        if ((res = nng_msg_alloc(&msg, sizeof(MsgMcstCommand))))
            dp_nng_fatal("ctrl/nng_msg_alloc", res);

        auto dpmsg = DP_NNG_BODY_AS(msg, MsgMcstCommand);
        dpmsg->head.msgType = DP_MT_McstCommand;
        dpmsg->cmd = DP_DC_RunDoom;

        if ((res = nng_sendmsg(ctx.pub, msg, 0)))
            dp_nng_fatal("ctrl/sendmsg", res);
    }
}

void final_cleanup(ControllerContext &ctx)
{
    log_debug("final cleanup: ending all dooms");
    end_all_dooms(ctx);
    std::this_thread::sleep_for(std::chrono::milliseconds(50));
    check_on_dooms(ctx);

    if (!ctx.dooms.empty())
    {
        log_warn("final cleanup: terminating all %zu remaining dooms", ctx.dooms.size());
        signal_all_dooms(ctx, SIGTERM);
    }
}

inline s32 RoundFloatToInt(float value)
{
    return static_cast<s32>(value + 0.5);
}

struct V4: public ImVec4
{
    using ImVec4::ImVec4;
    float &a = ImVec4::w;
    float &r = ImVec4::x;
    float &g = ImVec4::y;
    float &b = ImVec4::z;
};

inline u32 V4ToARGB(V4 c)
{
    u32 result = ((RoundFloatToInt(c.a * 255.0) & 0xff) << 24
                  | (RoundFloatToInt(c.r * 255.0) & 0xff) << 16
                  | (RoundFloatToInt(c.g * 255.0) & 0xff) <<  8
                  | (RoundFloatToInt(c.b * 255.0) & 0xff) <<  0);
    return result;
}

struct OffscreenBuffer
{
    s32 width;
    s32 height;
    u8 *pixels;
    s32 pitch;
    const int BytesPerPixel;
};

inline void PutPixelUnchecked(OffscreenBuffer *buffer, s32 x, s32 y, V4 color)
{
        u32 *pixel = reinterpret_cast<u32 *>(buffer->pixels
                                             + y * buffer->pitch
                                             + x * buffer->BytesPerPixel);
        *pixel = V4ToARGB(color);
}

inline void PutPixel(OffscreenBuffer *buffer, s32 x, s32 y, V4 color)
{
    if (x >= 0 && x < buffer->width && y >= 0 && y < buffer->height)
    {
        PutPixelUnchecked(buffer, x, y, color);
    }
}

void DrawRectangle(OffscreenBuffer *buffer,
        float realMinX, float realMinY,
        float realMaxX, float realMaxY,
        V4 color)

{
    s32 minX = RoundFloatToInt(realMinX);
    s32 minY = RoundFloatToInt(realMinY);
    s32 maxX = RoundFloatToInt(realMaxX);
    s32 maxY = RoundFloatToInt(realMaxY);

    if (minX < 0)
        minX = 0;

    if (minY < 0)
        minY = 0;

    if (maxX > buffer->width)
        maxX = buffer->width;

    if (maxY > buffer->height)
        maxY = buffer->height;


    u8 *row = buffer->pixels + minY * buffer->pitch;

    for (s32 y = minY; y < maxY; ++y)
    {
        u32 *dstPixel = reinterpret_cast<u32 *>(row + minX * buffer->BytesPerPixel);

        for (s32 x = minX; x < maxX; ++x)
        {
            float a = color.a;
            float srcR = color.r * 255.0;
            float srcG = color.g * 255.0;
            float srcB = color.b * 255.0;

            float dstR = ((*dstPixel >> 16) & 0xff);
            float dstG = ((*dstPixel >>  8) & 0xff);
            float dstB = ((*dstPixel >>  0) & 0xff);

            dstR = (1 - a) * dstR + a * srcR;
            dstG = (1 - a) * dstG + a * srcG;
            dstB = (1 - a) * dstB + a * srcB;

            *dstPixel = (RoundFloatToInt(dstR) << 16 |
                         RoundFloatToInt(dstG) <<  8 |
                         RoundFloatToInt(dstB) <<  0);

            ++dstPixel;
        }

        row += buffer->pitch;
    }
}

void render_dooms(ControllerContext &ctx)
{
    OffscreenBuffer buffer =
    {
        DoomScreenWidth,
        DoomScreenHeight,
        ctx.pixelBuffer.data(),
        DoomScreenWidth * DoomBytesPerPixel,
        DoomBytesPerPixel
    };

    DrawRectangle(&buffer, 0, 0, 10, 10, { 1, 0, 0, 1 }); // top-left red
    DrawRectangle(&buffer, 0, buffer.height-10, 10, buffer.height, { 0, 1, 0, 1 });
    DrawRectangle(&buffer, buffer.width-10, buffer.height-10, buffer.width, buffer.height, { 0, 0, 1, 1 });
    DrawRectangle(&buffer, buffer.width-10, 0, buffer.width, 10, { 0.840, 0.0168, 0.717, 1 });

    //DrawRectangle(&buffer, buffer.width-10, 0, buffer.width, 10, { 0.0, 1.0, 0.0, 1.0 }); // top-right green

    SDL_Rect destRect = {0, 0, buffer.width, buffer.height};
    const size_t doomCount = ctx.dooms.size();

    for (size_t i=0; i<doomCount; ++i)
    {
        if (i != 0 && i % ctx.columns == 0)
        {
            destRect.x = 0;
            destRect.y += buffer.height;
        }

        auto &ds = ctx.dooms.at(i);
        auto texture = ds.texture;
        assert(texture);
        //SDL_UpdateTexture(texture, nullptr, buffer.pixels, buffer.width * buffer.BytesPerPixel);
        SDL_RenderCopy(ctx.renderer, texture, nullptr, &destRect);

        destRect.x += buffer.width;
    }
}

ControllerActions run_ui(ControllerContext &ctx)
{
#ifndef IMGUI_DISABLE_DEBUG_TOOLS
    const bool has_debug_tools = true;
#else
    const bool has_debug_tools = false;
#endif
    static bool show_app_metrics = false;
    static bool show_app_debug_log = false;
    static bool show_log_window = true;

    if (show_app_metrics)
        ImGui::ShowMetricsWindow(&show_app_metrics);
    if (show_app_debug_log)
        ImGui::ShowDebugLogWindow(&show_app_debug_log);

    const ImGuiViewport* main_viewport = ImGui::GetMainViewport();

    if (show_log_window)
    {
        ImGui::SetNextWindowPos(ImVec2(main_viewport->WorkPos.x + 666, main_viewport->WorkPos.y + 20), ImGuiCond_FirstUseEver);
        ImGui::SetNextWindowSize(ImVec2(500, 400), ImGuiCond_FirstUseEver);
        ctx.appLog.Draw("log");
    }

    ImGui::SetNextWindowPos(ImVec2(main_viewport->WorkPos.x + 20, main_viewport->WorkPos.y + 20), ImGuiCond_FirstUseEver);
    ImGui::SetNextWindowSize(ImVec2(420, 340), ImGuiCond_FirstUseEver);
    ImGuiWindowFlags window_flags = ImGuiWindowFlags_MenuBar;

    static std::array<char, 1024> strbuf;
    auto &io = ImGui::GetIO();

    aprintf(strbuf, "doompanning - #dooms=%zu, %.2f ms/frame (%.1f fps)###doompanning",
        ctx.dooms.size(), 1000.0f / io.Framerate, io.Framerate);

    // Main body of the doompanning window starts here.
    if (!ImGui::Begin(strbuf.data(), nullptr, window_flags))
    {
        // Early out if the window is collapsed, as an optimization.
        ImGui::End();
        return {};
    }

    // Menu Bar
    if (ImGui::BeginMenuBar())
    {
        if (ImGui::BeginMenu("Menu"))
        {
            ImGui::MenuItem("Log Window", nullptr, &show_log_window);
            ImGui::MenuItem("Quit", "Ctrl+Q", &ctx.quit, true);
            ImGui::EndMenu();
        }

        if (ImGui::BeginMenu("Tools"))
        {
            ImGui::MenuItem("Dear ImGui Metrics/Debugger", NULL, &show_app_metrics, has_debug_tools);
            ImGui::MenuItem("Dear ImGui Debug Log", NULL, &show_app_debug_log, has_debug_tools);
            ImGui::EndMenu();
        }

        ImGui::EndMenuBar();
    }

    // Window contents

    ControllerActions result = {};
    static int doomsToSpawn = 1;

    ImGui::PushItemWidth(ImGui::GetFontSize() * -16); // affects stuff like slider widths

    ImGui::SliderInt("Layout columns##columns", &ctx.columns, 1, 32, "%d", ImGuiSliderFlags_AlwaysClamp);

    ImGui::SliderInt("##dooms", &doomsToSpawn, 1, 256, "%d", ImGuiSliderFlags_AlwaysClamp | ImGuiSliderFlags_Logarithmic);
    aprintf(strbuf, "Spawn %d more doom%s###spawnmore", doomsToSpawn, doomsToSpawn > 1 ? "s" : "");
    if (ImGui::SameLine(); ImGui::Button(strbuf.data()))
        result.doomsToSpawn = doomsToSpawn;

    if (ImGui::Button("End all Dooms"))
        result.endAllDooms = true;

    ImGui::PopItemWidth();
    ImGui::End();

    return result;
}

int doom_controller_loop(ControllerContext &ctx)
{
    static constexpr ImVec4 clear_color = ImVec4(0.45f, 0.55f, 0.60f, 1.00f);

    ControllerActions actions = {};

    while (!ctx.quit)
    {
        SDL_Event event;
        while (SDL_PollEvent(&event))
        {
            ImGui_ImplSDL2_ProcessEvent(&event);

            if (event.type == SDL_QUIT)
                ctx.quit = true;

            if (event.type == SDL_WINDOWEVENT
                && event.window.event == SDL_WINDOWEVENT_CLOSE
                && event.window.windowID == SDL_GetWindowID(ctx.window))
            {
                ctx.quit = true;
            }
        }

        // Process input events not consumed by ImGui
        if (auto &io = ImGui::GetIO();
            !io.WantCaptureKeyboard)
        {
            if (io.KeyCtrl && ImGui::IsKeyDown(ImGuiKey_Q))
            {
                ctx.quit = true;
            }
        }

        perform_actions(ctx, actions);
        check_on_dooms(ctx);
        do_networking(ctx);

        // Start the Dear ImGui frame
        ImGui_ImplSDLRenderer_NewFrame();
        ImGui_ImplSDL2_NewFrame();
        ImGui::NewFrame();

        actions = run_ui(ctx);

        // Rendering
        const auto [r, g, b, a] = imvec4_to_rgba(clear_color);
        SDL_SetRenderDrawColor(ctx.renderer, r, g, b, a);
        SDL_RenderClear(ctx.renderer);

        render_dooms(ctx);

        ImGui::Render();
        ImGui_ImplSDLRenderer_RenderDrawData(ImGui::GetDrawData());
        SDL_RenderPresent(ctx.renderer);
    }

    final_cleanup(ctx);

    return 0;
}

void log_to_imgui(log_Event *ev)
{
    auto ctx = reinterpret_cast<ControllerContext *>(ev->udata);
    ctx->appLog.AddLog(ev->fmt, ev->ap);
}

int main(int argc, char *argv[])
{
    (void) argc;
    (void) argv;

#ifndef NDEBUG
    log_set_level(LOG_TRACE);
#else
    log_set_level(LOG_DEBUG);
#endif

    log_info("doompanning ctrl starting");

    if (SDL_Init(SDL_INIT_VIDEO | SDL_INIT_TIMER))
        dp_sdl_fatal("SDL_Init");

#ifdef SDL_HINT_IME_SHOW_UI
    SDL_SetHint(SDL_HINT_IME_SHOW_UI, "1");
#endif

    const auto windowFlags = SDL_WINDOW_RESIZABLE | SDL_WINDOW_ALLOW_HIGHDPI | SDL_WINDOW_OPENGL;
    auto window = SDL_CreateWindow("doompanning", SDL_WINDOWPOS_CENTERED, SDL_WINDOWPOS_CENTERED, 1280, 720, windowFlags);
    if (!window)
        dp_sdl_fatal("SDL_CreateWindow");

    auto renderer = SDL_CreateRenderer(window, -1, SDL_RENDERER_PRESENTVSYNC | SDL_RENDERER_ACCELERATED);

    if (!renderer)
        dp_sdl_fatal("SDL_CreateRenderer");

    IMGUI_CHECKVERSION();
    ImGui::CreateContext();
    ImGui::GetIO().IniFilename = "doompanning_ui.ini";
    ImGui::StyleColorsDark();
    ImGui_ImplSDL2_InitForSDLRenderer(window, renderer);
    ImGui_ImplSDLRenderer_Init(renderer);

    dp_nng_init_limits(1, 1, 1); // int ncpu_max, int pool_thread_limit_max, int resolv_thread_limit

    ControllerContext ctx;
    ctx.pub = make_ctrl_pub(CtrlUrl);
    ctx.sub = make_ctrl_sub(DoomUrl);
    ctx.window = window;
    ctx.renderer = renderer;
    ctx.pixelBuffer.fill(0u);

    log_add_callback(log_to_imgui, &ctx, LOG_TRACE);

    int ret = doom_controller_loop(ctx);

    nng_close(ctx.pub);
    nng_close(ctx.sub);

    return ret;
}