#include "WindowGuard.h"
#include <unordered_map>

namespace {

struct GuardData {
    WNDPROC originalWndProc = nullptr;
    UINT flags = 0;
    bool allowClose = false;
    bool intentionalHide = false;
};

std::unordered_map<HWND, GuardData> g_guards;

const wchar_t* PROP_NAME = L"WindowGuardData";

LRESULT CALLBACK GuardWndProc(HWND hwnd, UINT msg, WPARAM wParam, LPARAM lParam) {
    auto it = g_guards.find(hwnd);
    if (it == g_guards.end()) {
        return DefWindowProc(hwnd, msg, wParam, lParam);
    }
    auto& gd = it->second;

    // 拦截关闭 (WM_CLOSE from DefWindowProc / DestroyWindow)
    if (msg == WM_CLOSE && (gd.flags & WindowGuard::BlockClose)) {
        if (!gd.allowClose) {
            return 0;  // Swallow the message
        }
    }

    // 拦截 SC_CLOSE (taskbar "Close window", Alt+F4, thumbnail X button).
    // WM_CLOSE above catches the DefWindowProc path, but many window
    // libraries (raylib/GLFW) handle SC_CLOSE directly before it ever
    // becomes WM_CLOSE.  Intercept it here so the taskbar can't kill us.
    if (msg == WM_SYSCOMMAND && wParam == SC_CLOSE && (gd.flags & WindowGuard::BlockClose)) {
        if (!gd.allowClose) {
            return 0;
        }
    }

    // 拦截最小化
    if (msg == WM_SIZE && wParam == SIZE_MINIMIZED && (gd.flags & WindowGuard::BlockMinimize)) {
        ShowWindow(hwnd, SW_RESTORE);
        return 0;
    }

    // 拦截外部隐藏
    if (msg == WM_WINDOWPOSCHANGING && (gd.flags & WindowGuard::BlockHide)) {
        auto* wp = reinterpret_cast<WINDOWPOS*>(lParam);
        if ((wp->flags & SWP_HIDEWINDOW) && !gd.intentionalHide) {
            wp->flags &= ~SWP_HIDEWINDOW;
            wp->flags |= SWP_SHOWWINDOW;
            return 0;
        }
        // 允许本次隐藏后复位标志
        if ((wp->flags & SWP_HIDEWINDOW) && gd.intentionalHide) {
            gd.intentionalHide = false;
        }
    }

    return CallWindowProc(gd.originalWndProc, hwnd, msg, wParam, lParam);
}

} // anonymous namespace

namespace WindowGuard {

WNDPROC Install(HWND hwnd, UINT flags) {
    if (g_guards.find(hwnd) != g_guards.end()) {
        return g_guards[hwnd].originalWndProc;
    }

    GuardData gd;
    gd.flags = flags;
    gd.originalWndProc = reinterpret_cast<WNDPROC>(
        SetWindowLongPtr(hwnd, GWLP_WNDPROC,
            reinterpret_cast<LONG_PTR>(GuardWndProc)));

    g_guards[hwnd] = gd;
    return gd.originalWndProc;
}

void Uninstall(HWND hwnd) {
    auto it = g_guards.find(hwnd);
    if (it == g_guards.end()) return;

    SetWindowLongPtr(hwnd, GWLP_WNDPROC,
        reinterpret_cast<LONG_PTR>(it->second.originalWndProc));
    g_guards.erase(it);
}

void SetAllowClose(HWND hwnd, bool allow) {
    auto it = g_guards.find(hwnd);
    if (it != g_guards.end()) {
        it->second.allowClose = allow;
    }
}

void SetIntentionalHide(HWND hwnd, bool intentional) {
    auto it = g_guards.find(hwnd);
    if (it != g_guards.end()) {
        it->second.intentionalHide = intentional;
    }
}

bool IsInstalled(HWND hwnd) {
    return g_guards.find(hwnd) != g_guards.end();
}

} // namespace WindowGuard
