#include "Console.h"
#include "../../../resource.h"
#include <cstdio>

std::atomic<bool> Console::s_bOutputPending{false};

// ---------------------------------------------------------------------------
// Singleton
// ---------------------------------------------------------------------------

Console* Console::Get()
{
    static Console instance;
    return &instance;
}

Console::~Console()
{
    Shutdown();
}

// ---------------------------------------------------------------------------
// Initialize / Shutdown
// ---------------------------------------------------------------------------

void Console::Initialize()
{
    if (m_bInitialized)
        return;
    m_bInitialized = true;

    // ---- Window identity ------------------------------------------------
    SetConsoleTitleW(L"Rcjy");
    m_originalTitle = L"Rcjy";

    m_hConWnd = GetConsoleWindow();
    if (m_hConWnd)
    {
        // Window icon
        HICON hIcon = LoadIconW(GetModuleHandleW(NULL), MAKEINTRESOURCEW(IDI_MAIN_ICON));
        if (hIcon)
            SendMessageW(m_hConWnd, WM_SETICON, ICON_BIG, (LPARAM)hIcon);

        HardenWindow(m_hConWnd);
    }

    // ---- stderr redirect ------------------------------------------------
    // Library output (libpng warnings etc.) goes to a log file so it never
    // blocks the console during QuickEdit selection.
    {
        wchar_t path[MAX_PATH];
        GetModuleFileNameW(NULL, path, MAX_PATH);
        wchar_t* p = wcsrchr(path, L'\\');
        if (p) *p = L'\0';
        wcscat_s(path, L"\\Logs\\stderr.log");
#pragma warning(push)
#pragma warning(disable : 4996)
        _wfreopen(path, L"a", stderr);
#pragma warning(pop)
    }

    // ---- Start guard thread ---------------------------------------------
    // Create the wake event BEFORE the thread so it is always valid.
    m_hWakeEvent = CreateEventW(NULL, FALSE, FALSE, NULL); // auto-reset
    m_bRunning   = true;
    m_hThread    = CreateThread(NULL, 0, GuardThreadProc, this, 0, NULL);

    // ---- Block Ctrl+C / Ctrl+Break -------------------------------------
    SetConsoleCtrlHandler(CtrlHandler, TRUE);
}

void Console::Shutdown()
{
    if (m_hThread == NULL)
        return;

    // Tell the thread to exit and wake it if it's sleeping.
    m_bRunning = false;
    if (m_hWakeEvent)
        SetEvent(m_hWakeEvent);

    WaitForSingleObject(m_hThread, 5000);
    CloseHandle(m_hThread);
    m_hThread = NULL;

    if (m_hWakeEvent)
    {
        CloseHandle(m_hWakeEvent);
        m_hWakeEvent = NULL;
    }
}

// ---------------------------------------------------------------------------
// Visibility
// ---------------------------------------------------------------------------

void Console::SetVisible(bool visible)
{
    // Refresh the handle in case the window was recreated.
    HWND hwnd = GetConsoleWindow();
    if (hwnd)
        m_hConWnd = hwnd;

    if (m_hConWnd)
        ShowWindow(m_hConWnd, visible ? SW_SHOW : SW_HIDE);
}

// ---------------------------------------------------------------------------
// One-shot selection cancel (called when menu opens)
// ---------------------------------------------------------------------------

void Console::CancelSelectionIfActive()
{
    // Refresh HWND in case it changed
    HWND hwnd = GetConsoleWindow();
    if (hwnd)
        m_hConWnd = hwnd;
    if (!m_hConWnd)
        return;

    wchar_t buf[256];
    if (GetWindowTextW(m_hConWnd, buf, 256) <= 0)
        return;

    // QuickEdit selection active?
    //   Chinese Windows : "选择 Rcjy"
    //   English Windows : "Select Rcjy"
    if (wcsstr(buf, L"选择") == buf || wcsstr(buf, L"Select") == buf)
    {
        PostMessageW(m_hConWnd, WM_KEYDOWN, VK_ESCAPE, 0x00010001);
        PostMessageW(m_hConWnd, WM_KEYUP,   VK_ESCAPE, 0xC0010001);
        SetConsoleTitleW(m_originalTitle.c_str());
    }
}

// ---------------------------------------------------------------------------
// Output-pending flags (called by Log)
// ---------------------------------------------------------------------------

void Console::NotifyOutputPending()
{
    s_bOutputPending.store(true, std::memory_order_release);
}

void Console::NotifyOutputDone()
{
    s_bOutputPending.store(false, std::memory_order_release);
}

// ---------------------------------------------------------------------------
// Ctrl+C / Ctrl+Break suppression
// ---------------------------------------------------------------------------
// Return TRUE to swallow the event so the console process doesn't terminate.
// CTRL_CLOSE_EVENT is blocked to prevent taskbar "Close window" from killing
// the process.  The console window may still be destroyed by conhost, but we
// survive and can be restarted from the overlay menu.

BOOL WINAPI Console::CtrlHandler(DWORD fdwCtrlType)
{
    switch (fdwCtrlType)
    {
    case CTRL_C_EVENT:
    case CTRL_BREAK_EVENT:
    case CTRL_CLOSE_EVENT:
        return TRUE;
    default:
        return FALSE;
    }
}

// ---------------------------------------------------------------------------
// Periodic maintenance
// ---------------------------------------------------------------------------

void Console::HardenWindow(HWND hwnd)
{
    // (1) Remove SC_CLOSE from the system menu → disables title-bar X,
    //     Alt+F4, and grays out taskbar "Close window" for most Windows
    //     versions.  Conhost is a separate process so we cannot subclass
    //     its WndProc; we can only influence the system menu.
    HMENU hMenu = GetSystemMenu(hwnd, FALSE);
    DeleteMenu(hMenu, SC_CLOSE, MF_BYCOMMAND);
    DrawMenuBar(hwnd);

    // (2) Add WS_EX_TOOLWINDOW to hide the console from the taskbar
    //     entirely.  No taskbar entry → no "Close window" to click.
    //     The window remains visible and accessible via Alt+Tab.
    LONG_PTR exStyle = GetWindowLongPtrW(hwnd, GWL_EXSTYLE);
    if (!(exStyle & WS_EX_TOOLWINDOW))
    {
        SetWindowLongPtrW(hwnd, GWL_EXSTYLE, exStyle | WS_EX_TOOLWINDOW);
        // Force the window to re-read its extended styles
        SetWindowPos(hwnd, NULL, 0, 0, 0, 0,
            SWP_NOACTIVATE | SWP_NOMOVE | SWP_NOSIZE | SWP_NOZORDER |
            SWP_NOOWNERZORDER | SWP_FRAMECHANGED);
    }
}

// ---------------------------------------------------------------------------
// Periodic maintenance
// ---------------------------------------------------------------------------

void Console::RefreshConsoleHandle()
{
    // Reacquire the console window handle.
    // conhost.exe may recreate the window in rare edge cases (e.g. after
    // the console is hidden/shown by an external tool).
    HWND hwnd = GetConsoleWindow();
    if (hwnd && hwnd != m_hConWnd)
    {
        m_hConWnd = hwnd;

        HICON hIcon = LoadIconW(GetModuleHandleW(NULL), MAKEINTRESOURCEW(IDI_MAIN_ICON));
        if (hIcon)
            SendMessageW(m_hConWnd, WM_SETICON, ICON_BIG, (LPARAM)hIcon);

        HardenWindow(m_hConWnd);
    }
    else if (hwnd)
    {
        m_hConWnd = hwnd;
        // Periodically re-apply hardening — belt and suspenders in case
        // something (e.g. a Windows update) reverts our changes.
        HardenWindow(m_hConWnd);
    }
    else if (m_hConWnd && !IsWindow(m_hConWnd))
    {
        // Console window was destroyed (e.g. external tool forcibly
        // closed it). The process survived because of CTRL_CLOSE_EVENT.
        m_hConWnd = NULL;
    }
}

// ---------------------------------------------------------------------------
// Title monitoring (called every 100 ms from the guard thread)
// ---------------------------------------------------------------------------
// When the user starts a QuickEdit selection, conhost changes the console
// title to "选择 <app>" (Chinese) or "Select <app>" (English).  We detect
// this and give the user a 500 ms grace period for short selections.  If
// the selection persists longer, we send Esc to cancel it — this prevents
// the selection from blocking console output before any WriteConsoleW call
// gets stuck.
// ---------------------------------------------------------------------------

void Console::CheckAndRestoreTitle()
{
    if (!m_hConWnd || m_originalTitle.empty())
        return;

    wchar_t buf[256];
    if (GetWindowTextW(m_hConWnd, buf, 256) <= 0)
        return;

    // QuickEdit selection title?
    //   Chinese Windows : "选择 Rcjy"
    //   English Windows : "Select Rcjy"
    if (wcsstr(buf, L"选择") == buf || wcsstr(buf, L"Select") == buf)
    {
        if (!m_bSelectionDetected)
        {
            m_bSelectionDetected = true;
            m_nSelectionDelay    = 5;   // 5 × 100ms = 500ms grace period
        }
        return; // Don't restore yet — let the countdown run
    }

    m_bSelectionDetected = false;

    // Any other external title change → restore immediately
    if (wcscmp(buf, m_originalTitle.c_str()) != 0)
        SetConsoleTitleW(m_originalTitle.c_str());
}

// ---------------------------------------------------------------------------
// Guard thread
// ---------------------------------------------------------------------------
// Two independent layers of QuickEdit protection:
//
// 1. TITLE-BASED (proactive) — conhost changes the title to "选择 Rcjy" /
//    "Select Rcjy" the moment the user enters selection mode.  We poll this
//    every 100 ms and give a 500 ms grace period for short selections before
//    cancelling with Esc.  This catches selection BEFORE any output blocks.
//
// 2. OUTPUT-BASED (reactive) — if selection still slips through and
//    WriteConsoleW gets stuck for >100 ms, we detect it via the
//    s_bOutputPending flag and send Esc.  This is the safety net.
//
// Both layers post Esc directly to conhost's window via PostMessageW, which
// bypasses the blocked console input queue because window messages are
// processed in conhost's own thread.
// ---------------------------------------------------------------------------

DWORD WINAPI Console::GuardThreadProc(LPVOID lpParam)
{
    Console* self = static_cast<Console*>(lpParam);
    if (!self)
        return 0;

    bool wasBlocked = false;
    int  tickCount  = 0;

    while (self->m_bRunning)
    {
        // Interruptible sleep — Shutdown() sets m_hWakeEvent to wake us.
        WaitForSingleObject(self->m_hWakeEvent, 100);

        if (!self->m_bRunning)
            break;

        // ---- Title check (every tick) -----------------------------------
        self->CheckAndRestoreTitle();

        // ---- Title-based selection countdown ----------------------------
        if (self->m_bSelectionDetected)
        {
            if (--self->m_nSelectionDelay <= 0)
            {
                // Only cancel the selection if output is actually blocked.
                // If nothing is waiting to write, let the user keep
                // selecting (e.g. to copy text).
                if (s_bOutputPending.load(std::memory_order_acquire) && self->m_hConWnd)
                {
                    // lParam encodes repeat-count=1, scan-code=0x01, extended=1
                    PostMessageW(self->m_hConWnd, WM_KEYDOWN, VK_ESCAPE, 0x00010001);
                    PostMessageW(self->m_hConWnd, WM_KEYUP,   VK_ESCAPE, 0xC0010001);
                    SetConsoleTitleW(self->m_originalTitle.c_str());
                    self->m_bSelectionDetected = false;
                }
                else
                {
                    // No output pending yet — reset the grace period
                    self->m_nSelectionDelay = 5;
                }
            }
        }

        // ---- Periodic HWND refresh (every ~500 ms) ----------------------
        if (++tickCount >= 5)
        {
            tickCount = 0;
            self->RefreshConsoleHandle();
        }

        // ---- Output-based blocking detection (reactive safety net) ------
        if (!s_bOutputPending.load(std::memory_order_acquire))
        {
            wasBlocked = false;
            continue;
        }

        // Output is in flight. Give it 100 ms to finish normally.
        WaitForSingleObject(self->m_hWakeEvent, 100);
        if (!self->m_bRunning)
            break;

        if (s_bOutputPending.load(std::memory_order_acquire))
        {
            // Still pending → WriteConsoleW is stuck in QuickEdit selection.
            // Post Esc to conhost's window to cancel the selection.
            if (!wasBlocked && self->m_hConWnd)
            {
                PostMessageW(self->m_hConWnd, WM_KEYDOWN, VK_ESCAPE, 0x00010001);
                PostMessageW(self->m_hConWnd, WM_KEYUP,   VK_ESCAPE, 0xC0010001);
                wasBlocked = true;
            }
        }
    }

    return 0;
}
