#include "tray-icon.h"
#include "win-platform.h"

#include <shellapi.h>
#include <iostream>

#define WM_TRAYICON (WM_USER + 1)
#define ID_TRAY_APP_ICON 1001
#define ID_TRAY_EXIT 1002
#define ID_TRAY_CONTROL_PANEL 1003

TrayIcon* TrayIcon::instance = nullptr;

TrayIcon::TrayIcon() : _hwnd(nullptr), _isConsoleVisible(true)
{
    instance = this;
}

TrayIcon::~TrayIcon()
{
    shutdown();
    instance = nullptr;
}

void* TrayIcon::windowProc(void* h, unsigned int msg, unsigned long long wParam, long long lParam)
{
    HWND hwnd = reinterpret_cast<HWND>(h);

    if (msg == WM_DESTROY)
    {
        PostQuitMessage(0);
        return 0;
    }

    if (msg == WM_TRAYICON)
    {
        if (lParam == WM_LBUTTONUP && instance)
            instance->onToggleConsole();
        else if (lParam == WM_RBUTTONUP)
        {
            POINT curPoint;
            GetCursorPos(&curPoint);

            HMENU popupHandle = CreatePopupMenu();
            InsertMenuA(popupHandle, 0, MF_BYPOSITION | MF_STRING, ID_TRAY_CONTROL_PANEL, "Open Control Panel");
            InsertMenuA(popupHandle, 1, MF_BYPOSITION | MF_STRING, ID_TRAY_EXIT, "Exit");

            SetForegroundWindow(hwnd);

            UINT clicked =
                TrackPopupMenu(popupHandle, TPM_RETURNCMD | TPM_NONOTIFY, curPoint.x, curPoint.y, 0, hwnd, NULL);

            if (clicked == ID_TRAY_CONTROL_PANEL && instance)
                instance->onOpenControlPanel();
            else if (clicked == ID_TRAY_EXIT && instance)
                instance->onExit();

            DestroyMenu(popupHandle);
        }
    }
    return reinterpret_cast<void*>(DefWindowProcA(hwnd, msg, wParam, lParam));
}

bool TrayIcon::init()
{
    WNDCLASSEXA wc = {};
    wc.cbSize = sizeof(WNDCLASSEX);
    wc.lpfnWndProc = reinterpret_cast<WNDPROC>(reinterpret_cast<void*>(windowProc));
    wc.hInstance = GetModuleHandle(NULL);
    wc.lpszClassName = "HexUnlockedTrayIconClass";

    if (!RegisterClassExA(&wc)) return false;

    HWND window = CreateWindowExA(0, "HexUnlockedTrayIconClass", "HexUnlocked", 0, 0, 0, 0, 0, HWND_MESSAGE, NULL,
                                  wc.hInstance, NULL);
    if (!window) return false;

    this->_hwnd = window;

    NOTIFYICONDATAA nid = {};
    nid.cbSize = sizeof(NOTIFYICONDATA);
    nid.hWnd = window;
    nid.uID = ID_TRAY_APP_ICON;
    nid.uFlags = NIF_ICON | NIF_MESSAGE | NIF_TIP;
    nid.uCallbackMessage = WM_TRAYICON;

    nid.hIcon = LoadIconA(GetModuleHandle(NULL), "IDI_ICON1");
    strncpy_s(nid.szTip, "HexUnlocked", sizeof(nid.szTip) - 1);
    nid.szTip[sizeof(nid.szTip) - 1] = '\0';

    Shell_NotifyIconA(NIM_ADD, &nid);

    return true;
}

void TrayIcon::shutdown()
{
    if (_hwnd)
    {
        NOTIFYICONDATAA nid = {};
        nid.cbSize = sizeof(NOTIFYICONDATA);
        nid.hWnd = reinterpret_cast<HWND>(_hwnd);
        nid.uID = ID_TRAY_APP_ICON;
        Shell_NotifyIconA(NIM_DELETE, &nid);

        DestroyWindow(reinterpret_cast<HWND>(_hwnd));
        UnregisterClassA("HexUnlockedTrayIconClass", GetModuleHandle(NULL));
        _hwnd = nullptr;
    }
}

void TrayIcon::processMessages()
{
    MSG msg;
    while (PeekMessage(&msg, NULL, 0, 0, PM_REMOVE))
    {
        if (msg.message == WM_QUIT)
        {
            break;
        }
        TranslateMessage(&msg);
        DispatchMessage(&msg);
    }

    HWND consoleWnd = GetConsoleWindow();
    if (consoleWnd && IsIconic(consoleWnd))
    {
        ShowWindow(consoleWnd, SW_HIDE);
        _isConsoleVisible = false;
    }
}

void TrayIcon::onToggleConsole()
{
    HWND consoleWnd = GetConsoleWindow();
    if (!consoleWnd) return;

    if (_isConsoleVisible)
    {
        ShowWindow(consoleWnd, SW_HIDE);
        _isConsoleVisible = false;
    }
    else
    {
        ShowWindow(consoleWnd, SW_RESTORE);

        HWND hCurWnd = GetForegroundWindow();
        DWORD dwMyID = GetCurrentThreadId();
        DWORD dwCurID = GetWindowThreadProcessId(hCurWnd, NULL);
        AttachThreadInput(dwCurID, dwMyID, TRUE);
        SetWindowPos(consoleWnd, HWND_TOPMOST, 0, 0, 0, 0, SWP_NOSIZE | SWP_NOMOVE);
        SetWindowPos(consoleWnd, HWND_NOTOPMOST, 0, 0, 0, 0, SWP_NOSIZE | SWP_NOMOVE);
        SetForegroundWindow(consoleWnd);
        SetFocus(consoleWnd);
        SetActiveWindow(consoleWnd);
        AttachThreadInput(dwCurID, dwMyID, FALSE);

        _isConsoleVisible = true;
    }
}

void TrayIcon::onOpenControlPanel()
{
    ShellExecuteA(NULL, "open", "https://dbd.neru.rip", NULL, NULL, SW_SHOWNORMAL);
}

void TrayIcon::onExit()
{
    if (_exitCallback) _exitCallback();
}

void TrayIcon::setExitCallback(std::function<void()> callback)
{
    _exitCallback = callback;
}
