#include "SysTray.h"

UINT SysTray::WM_TASKBARCREATED = RegisterWindowMessage(TEXT("TaskbarCreated"));

SysTray::SysTray(HWND hwnd, HICON hIcon, const wchar_t* tooltip)
    : m_hwnd(hwnd)
    , m_trayHwnd(nullptr)
    , m_hIcon(hIcon ? hIcon : LoadIcon(nullptr, IDI_APPLICATION))
    , m_timerId(0)
{
    // 生成唯一窗口类名
    wchar_t buf[64];
    swprintf_s(buf, L"SysTray_%p", this);
    m_className = buf;

    if (!RegisterWindowClass()) return;

    m_trayHwnd = CreateWindowEx(
        0, m_className.c_str(), L"",
        WS_OVERLAPPEDWINDOW,
        CW_USEDEFAULT, CW_USEDEFAULT,
        CW_USEDEFAULT, CW_USEDEFAULT,
        nullptr, nullptr, GetModuleHandle(nullptr), nullptr
    );
    if (!m_trayHwnd) return;

    // 将 this 指针存入窗口，窗口过程通过 GWLP_USERDATA 取回
    SetWindowLongPtr(m_trayHwnd, GWLP_USERDATA, reinterpret_cast<LONG_PTR>(this));

    ZeroMemory(&m_nid, sizeof(m_nid));
    m_nid.cbSize = sizeof(m_nid);
    m_nid.hWnd = m_trayHwnd;
    m_nid.uID = 1;
    m_nid.uFlags = NIF_MESSAGE | NIF_ICON | NIF_TIP;
    m_nid.uCallbackMessage = WM_TRAYICON;
    m_nid.hIcon = m_hIcon;
    wcscpy_s(m_nid.szTip, tooltip ? tooltip : L"");
}

SysTray::~SysTray() {
    StopAutoRefresh();
    RemoveIcon();
    if (m_trayHwnd) {
        SetWindowLongPtr(m_trayHwnd, GWLP_USERDATA, 0);
        DestroyWindow(m_trayHwnd);
        m_trayHwnd = nullptr;
    }
    UnregisterWindowClass();
}

bool SysTray::RegisterWindowClass() {
    WNDCLASSEX wc = {};
    wc.cbSize = sizeof(WNDCLASSEX);
    wc.lpfnWndProc = WindowProc;
    wc.hInstance = GetModuleHandle(nullptr);
    wc.lpszClassName = m_className.c_str();
    wc.hCursor = LoadCursor(nullptr, IDC_ARROW);
    wc.hbrBackground = static_cast<HBRUSH>(GetStockObject(WHITE_BRUSH));
    return RegisterClassEx(&wc) != 0;
}

void SysTray::UnregisterWindowClass() {
    UnregisterClass(m_className.c_str(), GetModuleHandle(nullptr));
}

bool SysTray::AddIcon() {
    if (!m_trayHwnd) return false;
    return Shell_NotifyIcon(NIM_ADD, &m_nid) != FALSE;
}

bool SysTray::RemoveIcon() {
    if (!m_trayHwnd) return false;
    return Shell_NotifyIcon(NIM_DELETE, &m_nid) != FALSE;
}

bool SysTray::RefreshIcon() {
    if (!m_trayHwnd) return false;
    m_nid.uFlags = NIF_ICON | NIF_TIP;
    return Shell_NotifyIcon(NIM_MODIFY, &m_nid) != FALSE;
}

void SysTray::ReaddIcon() {
    for (int retry = 0; retry < 3; ++retry) {
        RemoveIcon();  // best-effort cleanup; may fail if explorer already removed the icon
        if (AddIcon()) return;
        if (retry < 2) Sleep(100);
    }
}

void SysTray::StartAutoRefresh(UINT intervalMinutes) {
    if (m_timerId != 0) {
        KillTimer(m_trayHwnd, m_timerId);
    }
    if (intervalMinutes == 0) {
        m_timerId = 0;
        return;
    }
    m_timerId = SetTimer(m_trayHwnd, 1, intervalMinutes * 60 * 1000, nullptr);
}

void SysTray::StopAutoRefresh() {
    if (m_timerId != 0) {
        KillTimer(m_trayHwnd, m_timerId);
        m_timerId = 0;
    }
}

void SysTray::AddMenuItem(UINT id, const wchar_t* text, MenuCallback onSelect) {
    m_menuItems.push_back({id, text, std::move(onSelect)});
}

void SysTray::SetMenuItemChecked(UINT id, bool checked) {
    for (auto& item : m_menuItems) {
        if (item.id == id) {
            item.checked = checked;
            return;
        }
    }
}

void SysTray::ClearMenuItems() {
    m_menuItems.clear();
}

void SysTray::ShowContextMenu() {
    HMENU hMenu = CreatePopupMenu();
    if (!hMenu) return;

    for (const auto& item : m_menuItems) {
        UINT flags = MF_STRING;
        if (item.checked) flags |= MF_CHECKED;
        AppendMenu(hMenu, flags, item.id, item.text.c_str());
    }

    POINT pt;
    GetCursorPos(&pt);
    SetForegroundWindow(m_trayHwnd);
    TrackPopupMenu(hMenu, TPM_RIGHTBUTTON, pt.x, pt.y, 0, m_trayHwnd, nullptr);
    PostMessage(m_trayHwnd, WM_NULL, 0, 0);
    DestroyMenu(hMenu);
}

void SysTray::HandleMessage(UINT msg, WPARAM wParam, LPARAM lParam) {
    if (msg != WM_TRAYICON) return;

    switch (lParam) {
        case WM_RBUTTONUP:
            ShowContextMenu();
            break;
        case WM_LBUTTONUP:
            if (m_onLeftClick) m_onLeftClick();
            break;
        case WM_LBUTTONDBLCLK:
            if (m_onDoubleClick) m_onDoubleClick();
            break;
    }
}

LRESULT CALLBACK SysTray::WindowProc(HWND hwnd, UINT msg, WPARAM wParam, LPARAM lParam) {
    auto* self = reinterpret_cast<SysTray*>(GetWindowLongPtr(hwnd, GWLP_USERDATA));

    // 托盘图标消息
    if (self && msg == self->WM_TRAYICON) {
        self->HandleMessage(msg, wParam, lParam);
        return 0;
    }

    // Explorer 重启后重建托盘图标
    if (msg == WM_TASKBARCREATED) {
        if (self) self->ReaddIcon();
        return 0;
    }

    // 菜单命令
    if (msg == WM_COMMAND && self) {
        UINT id = LOWORD(wParam);
        // 先查菜单项回调
        for (const auto& item : self->m_menuItems) {
            if (item.id == id && item.callback) {
                item.callback(id);
                return 0;
            }
        }
        // 再查通用菜单回调
        if (self->m_menuCallback) {
            self->m_menuCallback(id);
        }
        return 0;
    }

    // 自动刷新定时器
    if (msg == WM_TIMER && self) {
        self->RefreshIcon();
        return 0;
    }

    return DefWindowProc(hwnd, msg, wParam, lParam);
}

// ============================================================================
// 便捷函数
// ============================================================================
SysTray* SysTrayUtil::CreateSimple(HWND hwnd, HICON hIcon, const wchar_t* tooltip) {
    auto* tray = new SysTray(hwnd, hIcon, tooltip);
    tray->SetOnDoubleClick([hwnd]() {
        ShowWindow(hwnd, SW_RESTORE);
        SetForegroundWindow(hwnd);
    });
    tray->AddMenuItem(999, L"退出");
    tray->AddIcon();
    tray->StartAutoRefresh(30);
    return tray;
}
