#include "impl_spice.h"

#include <windows.h>

#include "cfg/configurator.h"
#include "games/io.h"
#include "launcher/launcher.h"
#include "launcher/superexit.h"
#include "misc/eamuse.h"
#include "overlay/overlay.h"
#include "rawinput/rawinput.h"
#include "touch/touch.h"
#include "util/logging.h"
#include "util/time.h"
#include "util/utils.h"
#include "hooks/graphics/graphics.h"

#define DEBUG_VERBOSE 0

#if DEBUG_VERBOSE
#define log_debug(module, format_str, ...) logger::push( \
    LOG_FORMAT("M", module, format_str, ## __VA_ARGS__), logger::Style::GREY)
#else
#define log_debug(module, format_str, ...)
#endif

#if !defined(IMGUI_ENABLE_WIN32_DEFAULT_IME_FUNCTIONS) || \
    !defined(IMGUI_DISABLE_DEFAULT_ALLOCATORS) || \
    !defined(IMGUI_USE_BGRA_PACKED_COLOR) || \
    !defined(IMGUI_HAS_VIEWPORT) || \
    !defined(IMGUI_HAS_DOCK) || \
    !defined(IMGUI_DISABLE_DEMO_WINDOWS) || \
    defined(IMGUI_DISABLE_DEBUG_TOOLS)
#error "fix imconfig.h after updating imgui version"
#endif

// state
static HWND g_hWnd = nullptr;
static INT64 g_Time = 0;
static INT64 g_TicksPerSecond = 0;
static ImGuiMouseCursor g_LastMouseCursor = ImGuiMouseCursor_COUNT;
static double g_LastMouseMovement = 0.f;
static bool g_MouseCursorAutoHide = false;
static float g_DisplaySizeOverrideW = 0.0f;
static float g_DisplaySizeOverrideH = 0.0f;

constexpr size_t VKEY_MAX = 255;
static std::array<bool, VKEY_MAX> g_KeysDown;
static std::array<bool, ImGuiMouseButton_COUNT> g_MouseDown;

static ImGuiKey get_imgui_key(int vkey) {
    switch (vkey) {
        case VK_TAB:
            return ImGuiKey_Tab;
        case VK_LEFT:
            return ImGuiKey_LeftArrow; 
        case VK_RIGHT:
            return ImGuiKey_RightArrow;
        case VK_UP:
            return ImGuiKey_UpArrow;
        case VK_DOWN:
            return ImGuiKey_DownArrow;
        case VK_PRIOR:
            return ImGuiKey_PageUp;
        case VK_NEXT:
            return ImGuiKey_PageDown;
        case VK_HOME:
            return ImGuiKey_Home;
        case VK_END:
            return ImGuiKey_End;
        case VK_INSERT:
            return ImGuiKey_Insert;
        case VK_DELETE:
            return ImGuiKey_Delete;
        case VK_BACK:
            return ImGuiKey_Backspace;
        case VK_SPACE:
            return ImGuiKey_Space;
        case VK_RETURN:
            return ImGuiKey_Enter;
        case VK_ESCAPE:
            return ImGuiKey_Escape;
        case VK_SHIFT:
        case VK_LSHIFT:
            return ImGuiKey_LeftShift;
        case VK_RSHIFT:
            return ImGuiKey_RightShift;
        case VK_CONTROL:
        case VK_LCONTROL:
            return ImGuiKey_LeftCtrl;
        case VK_RCONTROL:
            return ImGuiKey_RightCtrl;
        case 'A':
        case 'a':
            return ImGuiKey_A;
        case 'C':
        case 'c':
            return ImGuiKey_C;
        case 'V':
        case 'v':
            return ImGuiKey_V;
        case 'X':
        case 'x':
            return ImGuiKey_X;
        case 'Y':
        case 'y':
            return ImGuiKey_Y;  
        case 'Z':
        case 'z':
            return ImGuiKey_Z;
        default:
            return ImGuiKey_None;
    }
}

static ImGuiKey get_imgui_mod_key(int vkey) {
    switch (vkey) {
        case VK_SHIFT:
        case VK_LSHIFT:
        case VK_RSHIFT:
            return ImGuiMod_Shift;
            return ImGuiMod_Shift;
        case VK_CONTROL:
        case VK_LCONTROL:
        case VK_RCONTROL:
            return ImGuiMod_Ctrl;
        case VK_MENU:
        case VK_LMENU:
        case VK_RMENU:
            return ImGuiMod_Alt;
        default:
            return ImGuiMod_None;
    }
}

bool ImGui_ImplSpice_Init(HWND hWnd) {
    log_misc("imgui_impl_spice", "init");

    // check if already initialized
    if (g_hWnd != nullptr) {
        if (g_hWnd == hWnd) {
            return true;
        } else {
            ImGui_ImplSpice_Shutdown();
        }
    }

    // init performance stuff
    if (!::QueryPerformanceFrequency((LARGE_INTEGER *)&g_TicksPerSecond))
        return false;
    if (!::QueryPerformanceCounter((LARGE_INTEGER *)&g_Time))
        return false;

    // setup back-end capabilities flags
    g_hWnd = hWnd;
    ImGuiIO &io = ImGui::GetIO();
    io.BackendFlags |= ImGuiBackendFlags_HasMouseCursors;
    io.BackendFlags |= ImGuiBackendFlags_HasSetMousePos;
    io.BackendPlatformName = "imgui_impl_spice";
    io.ConfigErrorRecoveryEnableTooltip = true;

    // get display size
    ImGui_ImplSpice_UpdateDisplaySize();

    // only enable mouse cursor auto hiding if we want ImGui to draw the cursor
    // (if OS is drawing it, we don't auto-hide)
    if (io.MouseDrawCursor && !cfg::CONFIGURATOR_STANDALONE) {
        g_MouseCursorAutoHide = true;
    }

    // return success
    return true;
}

void ImGui_ImplSpice_Shutdown() {
    log_misc("imgui_impl_spice", "shutdown");

    // reset window handle
    g_hWnd = nullptr;
}

void ImGui_ImplSpice_UpdateDisplaySize() {

    // backends that render to a surface different from the window's client
    // area (e.g. dx11 swapchain backbuffer) can override the size here so
    // ImGui's viewport matches the actual render target.
    if (g_DisplaySizeOverrideW > 0.0f && g_DisplaySizeOverrideH > 0.0f) {
        ImGui::GetIO().DisplaySize = ImVec2(g_DisplaySizeOverrideW, g_DisplaySizeOverrideH);
        return;
    }

    // get display size
    RECT rect;
    ::GetClientRect(g_hWnd, &rect);
    ImGui::GetIO().DisplaySize = ImVec2((float)(rect.right - rect.left), (float)(rect.bottom - rect.top));
}

void ImGui_ImplSpice_SetDisplaySizeOverride(float w, float h) {
    g_DisplaySizeOverrideW = w;
    g_DisplaySizeOverrideH = h;
}

bool ImGui_ImplSpice_UpdateMouseCursor() {

    // check if cursor should be changed
    auto &io = ImGui::GetIO();
    if (io.ConfigFlags & ImGuiConfigFlags_NoMouseCursorChange) {
        return false;
    }

    // update cursor
    ImGuiMouseCursor imgui_cursor = ImGui::GetMouseCursor();

    // in auto-hide mode imgui owns cursor drawing, so the OS cursor must stay
    // hidden. this function is also called from the game window's WM_SETCURSOR
    // handler, which fires on every mouse move; without forcing it hidden here,
    // the else branch below would set IDC_ARROW on the next mouse move once the
    // overlay is closed. that arrow is conspicuous on games whose window class
    // has hCursor==NULL (e.g. DDR), since those normally show no client-area
    // cursor at all.
    if (imgui_cursor == ImGuiMouseCursor_None || io.MouseDrawCursor || g_MouseCursorAutoHide) {

        // hide OS mouse cursor if imgui is drawing it or if it wants no cursor
        ::SetCursor(nullptr);

    } else {

        // show OS mouse cursor
        LPTSTR win32_cursor = IDC_ARROW;
        switch (imgui_cursor) {
            case ImGuiMouseCursor_Arrow:
                win32_cursor = IDC_ARROW;
                break;
            case ImGuiMouseCursor_TextInput:
                win32_cursor = IDC_IBEAM;
                break;
            case ImGuiMouseCursor_ResizeAll:
                win32_cursor = IDC_SIZEALL;
                break;
            case ImGuiMouseCursor_ResizeEW:
                win32_cursor = IDC_SIZEWE;
                break;
            case ImGuiMouseCursor_ResizeNS:
                win32_cursor = IDC_SIZENS;
                break;
            case ImGuiMouseCursor_ResizeNESW:
                win32_cursor = IDC_SIZENESW;
                break;
            case ImGuiMouseCursor_ResizeNWSE:
                win32_cursor = IDC_SIZENWSE;
                break;
            case ImGuiMouseCursor_Hand:
                win32_cursor = IDC_HAND;
                break;
            default:
                break;
        }
        ::SetCursor(::LoadCursor(nullptr, win32_cursor));
    }
    return true;
}

static void ImGui_ImplSpice_UpdateMousePos() {

    // get current window size
    RECT rect;
    if (GetClientRect(g_hWnd, &rect)) {
        ImVec2 window_size(
                (float)(rect.right - rect.left),
                (float)(rect.bottom - rect.top));

        // set OS mouse position if requested
        auto &io = ImGui::GetIO();
        if (io.WantSetMousePos) {
            POINT pos {
                .x = static_cast<long>(io.MousePos.x),
                .y = static_cast<long>(io.MousePos.y),
            };
            ::ClientToScreen(g_hWnd, &pos);
            ::SetCursorPos(
                    static_cast<int>(pos.x / io.DisplaySize.x * window_size.x),
                    static_cast<int>(pos.y / io.DisplaySize.y * window_size.y));
        }

        const auto active_window = ::GetForegroundWindow();

        // if the main focus is a windowed subscreen, put the imgui cursor in a place that won't
        // trigger any overlay, don't process anything else
        const auto is_windowed_subscreen = 
            (GRAPHICS_IIDX_WSUB && active_window == TDJ_SUBSCREEN_WINDOW) ||
            (active_window == SDVX_SUBSCREEN_WINDOW) ||
            (active_window == POPN_SUBSCREEN_WINDOW);
        if (is_windowed_subscreen) {
            io.MousePos = ImVec2(-FLT_MAX, -FLT_MAX);
            return;
        }

        // set mouse position
        io.MousePos = ImVec2(-FLT_MAX, -FLT_MAX);
        POINT pos;
        if (active_window) {
            if (active_window == g_hWnd
            || ::IsChild(active_window, g_hWnd)
            || ::IsChild(g_hWnd, active_window)
            || active_window == SPICETOUCH_TOUCH_HWND) {
                if (::GetCursorPos(&pos) && ::ScreenToClient(g_hWnd, &pos)) {
                    io.MousePos = ImVec2(
                            (float) pos.x * io.DisplaySize.x / window_size.x,
                            (float) pos.y * io.DisplaySize.y / window_size.y);
                }
            }
        }

        // fallback to touch hwnd
        if (io.MousePos.x == -FLT_MAX || io.MousePos.y == -FLT_MAX) {
            if (SPICETOUCH_TOUCH_HWND) {
                if (::GetCursorPos(&pos) && ::ScreenToClient(SPICETOUCH_TOUCH_HWND, &pos)) {
                    io.MousePos = ImVec2(
                            (float) pos.x * io.DisplaySize.x / window_size.x,
                            (float) pos.y * io.DisplaySize.y / window_size.y);
                }
            }
        }

        // alternatively check touch
        std::vector<TouchPoint> touch_points;
        touch_get_points(touch_points, true);
        static size_t delay_touch = 0;
        static size_t delay_touch_target = 2;
        static DWORD last_touch_id = ~0u;
        if (!touch_points.empty() && !is_windowed_subscreen) {

            // use the first touch point
            auto &tp = touch_points[0];
            io.MousePos.x = tp.x * io.DisplaySize.x / window_size.x;
            io.MousePos.y = tp.y * io.DisplaySize.y / window_size.y;

            // update cursor position
            if (!tp.mouse) {
                pos.x = static_cast<long>(io.MousePos.x);
                pos.y = static_cast<long>(io.MousePos.y);
                ::ClientToScreen(g_hWnd, &pos);
                ::SetCursorPos(
                        static_cast<long>(pos.x / io.DisplaySize.x * window_size.x),
                        static_cast<long>(pos.y / io.DisplaySize.y * window_size.y));
            }

            // delay press
            g_MouseDown[ImGuiMouseButton_Left] = delay_touch++ >= delay_touch_target && last_touch_id == tp.id;
            if (last_touch_id == ~0u) {
                last_touch_id = tp.id;
            }

        } else {

            // reset
            delay_touch = 0;
            last_touch_id = ~0;
        }
    }
}

// Decide whether ImGui should draw its software cursor this frame. Kept out of
// ImGui_ImplSpice_NewFrame so the visibility rules live in one place.
//
// Only matters when ImGui owns cursor drawing (g_MouseCursorAutoHide, i.e. the
// game hid the OS cursor); otherwise the OS draws the cursor and
// io.MouseDrawCursor stays false throughout. In that mode:
//   - overlay hidden -> never draw the cursor (input belongs to the game)
//   - overlay shown  -> visible immediately, auto-hide after 2s idle, reappear
//                       on any mouse movement / wheel / button
static void ImGui_ImplSpice_UpdateCursorVisibility(
        bool overlay_visible,
        const ImVec2 &old_mouse_pos,
        const ImVec2 &new_mouse_pos,
        long wheel_diff)
{
    if (!g_MouseCursorAutoHide) {
        return;
    }

    auto &io = ImGui::GetIO();
    static bool overlay_visible_old = false;

    const bool mouse_activity =
        old_mouse_pos.x != new_mouse_pos.x ||
        old_mouse_pos.y != new_mouse_pos.y ||
        wheel_diff != 0 ||
        g_MouseDown[ImGuiMouseButton_Left] ||
        g_MouseDown[ImGuiMouseButton_Right] ||
        g_MouseDown[ImGuiMouseButton_Middle];

    if (!overlay_visible) {
        // overlay hidden: the game owns the cursor, don't draw ImGui's
        io.MouseDrawCursor = false;
    } else if (!overlay_visible_old || mouse_activity) {
        // overlay just opened, or the mouse moved: show the cursor and
        // (re)start the idle timer
        g_LastMouseMovement = get_performance_milliseconds();
        io.MouseDrawCursor = true;
    } else if ((get_performance_milliseconds() - g_LastMouseMovement) > 2000) {
        // mouse idle for more than 2 seconds while the overlay is open: hide it
        io.MouseDrawCursor = false;
    }

    overlay_visible_old = overlay_visible;
}

void ImGui_ImplSpice_NewFrame() {

    // check if font is built
    ImGuiIO& io = ImGui::GetIO();
    IM_ASSERT(io.Fonts->IsBuilt());

    // setup time step
    INT64 current_time;
    ::QueryPerformanceCounter((LARGE_INTEGER *)&current_time);
    io.DeltaTime = (float) (current_time - g_Time) / g_TicksPerSecond;
    g_Time = current_time;

    // only process new input from devices if window is in focus
    // (when not in focus, all mouse/keyboard events are treated as released)
    const auto overlay_visible = overlay::OVERLAY && overlay::OVERLAY->get_active();
    const auto accept_new_input = superexit::has_focus() && !rawinput::OS_WINDOW_ACTIVE && overlay_visible;

    // when running as standalone spicecfg.exe the configurator window proc feeds
    // ImGui directly via WM_KEY*/WM_MOUSE*/WM_MOUSEWHEEL messages. The per-frame
    // rawinput device walk and the 256-VK scan below are pure CPU waste in that
    // case (and the dominant per-frame cost on low-end PCs), so short-circuit
    // them entirely. Rebind dialogs that explicitly need fresh device state call
    // RI_MGR->devices_get_updated() themselves; see overlay/windows/config.cpp.
    const bool drive_input_from_rawinput = !cfg::CONFIGURATOR_STANDALONE;

    // remember old state
    std::array<BYTE, VKEY_MAX> KeysDownOld;
    for (size_t i = 0; i < VKEY_MAX; i++) {
        KeysDownOld[i] = g_KeysDown[i] ? ~0 : 0;
    }

    const auto MouseDownOld = g_MouseDown;

    // reset keys state
    if (drive_input_from_rawinput) {
        g_MouseDown.fill(false);
        g_KeysDown.fill(false);
    }

    // apply windows mouse buttons
    if (accept_new_input && drive_input_from_rawinput) {
        g_MouseDown[ImGuiMouseButton_Left] |= get_async_primary_mouse();
        g_MouseDown[ImGuiMouseButton_Right] |= get_async_secondary_mouse();
        g_MouseDown[ImGuiMouseButton_Middle] |= (GetAsyncKeyState(VK_MBUTTON) & 0x8000) != 0;
    }

    // read new keys state
    static long mouse_wheel_last = 0;
    long mouse_wheel = 0;
    if (drive_input_from_rawinput && RI_MGR != nullptr) {
        auto devices = RI_MGR->devices_get();
        for (auto &device : devices) {
            switch (device.type) {
                case rawinput::MOUSE: {
                    auto &mouse = device.mouseInfo;

                    // mouse button triggers
                    if (accept_new_input) {
                        if (GetSystemMetrics(SM_SWAPBUTTON)) {
                            if (mouse->key_states[rawinput::MOUSEBTN_RIGHT]) {
                                g_MouseDown[ImGuiMouseButton_Left] = true;
                            }
                            if (mouse->key_states[rawinput::MOUSEBTN_LEFT]) {
                                g_MouseDown[ImGuiMouseButton_Right] = true;
                            }
                        } else {
                            if (mouse->key_states[rawinput::MOUSEBTN_LEFT]) {
                                g_MouseDown[ImGuiMouseButton_Left] = true;
                            }
                            if (mouse->key_states[rawinput::MOUSEBTN_RIGHT]) {
                                g_MouseDown[ImGuiMouseButton_Right] = true;
                            }
                        }
                        if (mouse->key_states[rawinput::MOUSEBTN_MIDDLE]) {
                            g_MouseDown[ImGuiMouseButton_Middle] = true;
                        }
                    }

                    // final mouse wheel value should be all devices combined
                    // need to continuously calculate mouse wheel position even when not in focus
                    // to avoid values jumping (since mouse wheel values are absolute not relative)
                    mouse_wheel += mouse->pos_wheel;

                    break;
                }
                case rawinput::KEYBOARD: {
                    if (accept_new_input) {
                        // iterate all virtual key codes
                        for (size_t vKey = 0; vKey < VKEY_MAX; vKey++) {
                            // get state (combined from all pages)
                            auto &key_states = device.keyboardInfo->key_states;
                            for (size_t page_index = 0; page_index < 1024; page_index += 256) {
                                g_KeysDown[vKey] |= key_states[page_index + vKey];
                            }
                        }
                    }

                    break;
                }
                default:
                    break;
            }
        }
    }

    // process keyboard input from all keyboards collapsed into one state (g_KeysDown).
    // Skipped entirely in standalone configurator mode where Win32 messages already
    // drive io.AddKeyEvent / io.AddInputCharacter directly.
    if (drive_input_from_rawinput) {
        for (size_t vKey = 0; vKey < VKEY_MAX; vKey++) {
            const bool state = g_KeysDown[vKey];
            const auto imgui_key = get_imgui_key(vKey);
            const auto changed =
                (state && !KeysDownOld[vKey]) ||
                (!state && KeysDownOld[vKey]);

            if (imgui_key != ImGuiKey_None && changed) {
                io.AddKeyEvent(imgui_key, state);
                log_debug("imgui_impl_spice", "vkey {:#x} added as navigation event, state: {}", static_cast<uint64_t>(vKey), state);

                // mod key must also be processed separately
                const auto imgui_mod_key = get_imgui_mod_key(vKey);
                if (imgui_mod_key != ImGuiMod_None) {
                    io.AddKeyEvent(imgui_mod_key, state);
                }
            }

            // generate character input, but only if WM_CHAR didn't take over the functionality
            // only detecting rising edges here - this means holding a key won't work
            // (it's better than repeating a character input every frame - cost we pay for providing input
            // on top of rawinput instead of WM_CHAR)
            if (!overlay::USE_WM_CHAR_FOR_IMGUI_CHAR_INPUT && !KeysDownOld[vKey] && state) {
                UCHAR buf[2];
                auto ret = ToAscii(
                        static_cast<UINT>(vKey),
                        0,
                        static_cast<const BYTE *>(KeysDownOld.data()),
                        reinterpret_cast<LPWORD>(buf),
                        0);
                if (ret > 0) {
                    for (int i = 0; i < ret; i++) {
                        overlay::OVERLAY->input_char(buf[i]);
                        log_debug("imgui_impl_spice", "vkey {:#x} inputted as character", vKey);
                    }
                }
            }
        }
    }

    // set mouse wheel
    long wheel_diff = mouse_wheel - mouse_wheel_last;
    mouse_wheel_last = mouse_wheel;
    if (wheel_diff != 0 && accept_new_input && drive_input_from_rawinput) {
        io.AddMouseWheelEvent(0, wheel_diff);
    }

    // update OS mouse position. The standalone configurator gets mouse position
    // straight from WM_MOUSEMOVE, so skip the per-frame cursor poll there.
    const auto old_mouse_pos = io.MousePos;
    if (accept_new_input && drive_input_from_rawinput) {
        ImGui_ImplSpice_UpdateMousePos();
    }
    const auto new_mouse_pos = io.MousePos;

    // update mouse buttons
    // doing this after ImGui_ImplSpice_UpdateMousePos since it can set g_MouseDown for touch input
    if (drive_input_from_rawinput) {
        for (size_t i = 0; i < g_MouseDown.size(); i++) {
            if (MouseDownOld[i] != g_MouseDown[i]) {
                io.AddMouseButtonEvent(i, g_MouseDown[i]);
                log_debug("imgui_impl_spice", "mouse button {} event", g_MouseDown[i]);
            }
        }
    }

    // mouse cursor: visibility (overlay visibility + auto-hide)
    ImGui_ImplSpice_UpdateCursorVisibility(overlay_visible, old_mouse_pos, new_mouse_pos, wheel_diff);

    if (cfg::CONFIGURATOR_STANDALONE) {
        // if cursor is inside the client area, always set the OS cursor to what ImGui wants
        // this is to deal with cases where mouse cursor changes outside the client rect and comes
        // back into the window
        // i'm sure there might be better ways to deal with this but this works so whatever, right?
        RECT client_rect;
        if (GetClientRect(g_hWnd, &client_rect)) {
            POINT cursor;
            if (GetCursorPos(&cursor) && ScreenToClient(g_hWnd, &cursor)) {
                if (client_rect.left < cursor.x && cursor.x < client_rect.right &&
                    client_rect.top < cursor.y && cursor.y < client_rect.bottom) {
                    ImGui_ImplSpice_UpdateMouseCursor();
                }
            }
        }
    } else {
        // in auto-hide mode imgui owns the cursor, so keep the OS cursor hidden
        // even when the overlay is closed (io.MouseDrawCursor is false then).
        ImGuiMouseCursor mouse_cursor =
            (io.MouseDrawCursor || g_MouseCursorAutoHide)
                ? ImGuiMouseCursor_None
                : ImGui::GetMouseCursor();
        if (g_LastMouseCursor != mouse_cursor) {
            g_LastMouseCursor = mouse_cursor;
            ImGui_ImplSpice_UpdateMouseCursor();
        }
    }
}
