#include "steamvr_watch.h"
#include "hid_handler.h"
#include "ui.h"
#include <TlHelp32.h>
#include <algorithm>
#include <cwctype>
#include <iostream>


SteamVRWatch steamvrWatch;

namespace
{
    struct EnumWindowContext
    {
        DWORD pid;
        bool sentClose = false;
    };

    BOOL CALLBACK CloseWindowForPid(HWND hwnd, LPARAM lParam)
    {
        auto* context = reinterpret_cast<EnumWindowContext*>(lParam);
        DWORD windowPid = 0;
        GetWindowThreadProcessId(hwnd, &windowPid);

        if (windowPid == context->pid && IsWindowVisible(hwnd))
        {
            PostMessageW(hwnd, WM_CLOSE, 0, 0);
            context->sentClose = true;
        }

        return TRUE;
    }

    std::wstring NormalizeProcessName(std::wstring value)
    {
        std::transform(value.begin(), value.end(), value.begin(),
            [](wchar_t c) { return static_cast<wchar_t>(std::towlower(c)); });

        constexpr wchar_t kExeSuffix[] = L".exe";
        if (value.size() > 4 && value.compare(value.size() - 4, 4, kExeSuffix) == 0)
        {
            value.resize(value.size() - 4);
        }
        return value;
    }

    std::vector<DWORD> FindProcessIdsByName(const std::wstring& processNameNoExtension)
    {
        std::vector<DWORD> ids;
        const std::wstring target = NormalizeProcessName(processNameNoExtension);

        HANDLE snapshot = CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, 0);
        if (snapshot == INVALID_HANDLE_VALUE)
        {
            return ids;
        }

        PROCESSENTRY32W pe{};
        pe.dwSize = sizeof(pe);

        if (Process32FirstW(snapshot, &pe))
        {
            while (Process32NextW(snapshot, &pe))
            {
                if (NormalizeProcessName(pe.szExeFile) == target)
                {
                    ids.push_back(pe.th32ProcessID);
                }
            }
        }

        CloseHandle(snapshot);
        return ids;
    }

    bool RequestCloseThenKill(DWORD pid, DWORD waitMs)
    {
        HANDLE process = OpenProcess(PROCESS_TERMINATE | SYNCHRONIZE, FALSE, pid);
        if (!process)
        {
            return false;
        }

        EnumWindowContext context{ pid, false };
        EnumWindows(CloseWindowForPid, reinterpret_cast<LPARAM>(&context));

        if (context.sentClose)
        {
            const DWORD waitResult = WaitForSingleObject(process, waitMs);
            if (waitResult == WAIT_OBJECT_0)
            {
                CloseHandle(process);
                return true;
            }
        }

        const BOOL terminated = TerminateProcess(process, 0);
        if (terminated)
        {
            WaitForSingleObject(process, waitMs);
        }

        CloseHandle(process);
        return terminated == TRUE;
    }

    bool ProcessHasVisibleWindow(DWORD pid, std::wstring& outTitle)
    {
        struct EnumData
        {
            DWORD pid;
            std::wstring* title;
            bool found = false;
        };

        EnumData data{ pid, &outTitle, false };

        EnumWindows(
            [](HWND hwnd, LPARAM lParam) -> BOOL
            {
                auto* data = reinterpret_cast<EnumData*>(lParam);

                if (!IsWindowVisible(hwnd))
                    return TRUE;

                if (GetWindow(hwnd, GW_OWNER) != nullptr)
                    return TRUE;

                DWORD windowPid = 0;
                GetWindowThreadProcessId(hwnd, &windowPid);

                if (windowPid != data->pid)
                    return TRUE;

                wchar_t title[256];
                GetWindowTextW(hwnd, title, 256);

                if (wcslen(title) == 0)
                    return TRUE;

                data->found = true;
                *data->title = title;

                return FALSE;
            },
            reinterpret_cast<LPARAM>(&data));

        return data.found;
    }

}

bool SteamVRWatch::BlockFromSteamVRAbsent()
{
    bool runUnderApp = ui.get_runtime_state()->et_selected_apps.size() > 0;

    return ui.get_runtime_state()->et_require_steamvr_presence &&
        ((runUnderApp && (!steamvrActive.load() || !selectedAppsActive.load())) || (!runUnderApp && !steamvrActive.load()));
}

void SteamVRWatch::StartWatch()
{
    if (watchRunning.load())
    {
        return;
    }

    hid_handler.log << "Starting SteamVR watcher" << std::endl;
    watchRunning.store(true);

    // Create or reset the wake event
    if (wakeEvent_ == nullptr)
    {
        wakeEvent_ = CreateEventW(nullptr, FALSE, FALSE, nullptr);
        if (wakeEvent_ == nullptr)
        {
            hid_handler.log << "Failed to create wake event. Win32 error: " << GetLastError() << std::endl;
            return;
        }
    }
    else
    {
        ResetEvent(wakeEvent_);
    }

    stopRequested_.store(false);
    steamVrMonitor_ = std::thread([this]()
    {
        FindSteamVR(); 
    });
}

void SteamVRWatch::EndWatch()
{
    stopRequested_.store(true);
    watchRunning.store(false);

    // Signal the event to wake up the thread immediately
    if (wakeEvent_ != nullptr)
    {
        SetEvent(wakeEvent_);
    }

    if (steamVrMonitor_.joinable())
    {
        hid_handler.log << "Ending SteamVR watcher" << std::endl;
        steamVrMonitor_.join();
    }

    CloseAlignmentProcess();
}

void SteamVRWatch::FindSteamVR()
{
    while (!stopRequested_.load())
    {
        try
        {
            // See if the user's selected app is running
            bool found_app = false;
            for (RunningApp a : ui.get_runtime_state()->et_selected_apps)
            {
                if (IsProcessRunningByName(a.processName))
                {
                    found_app = true;
                }
            }

            selectedAppsActive.store(found_app);
            steamvrActive.store(IsProcessRunningByName(kTargetProcess));
            const bool helperRunning = IsHelperRunning();

            if (ui.get_runtime_state()->enabled_alignment_helper && steamvrActive.load() && !BlockFromSteamVRAbsent() && !helperRunning)
            {
                StartAlignmentProcess();
            }
            else if ((!ui.get_runtime_state()->enabled_alignment_helper || !steamvrActive.load() || BlockFromSteamVRAbsent()) && helperRunning)
            {
                CloseAlignmentProcess();
            }
        }
        catch (...)
        {
            hid_handler.log << "Error in SteamVR watch loop" << std::endl;
        }

        // Wait for 6 seconds or until woken by end
        if (wakeEvent_ != nullptr)
        {
            WaitForSingleObject(wakeEvent_, loop_delay);
        }
        else
        {
            Sleep(loop_delay);
        }
    }
}

bool SteamVRWatch::StartAlignmentProcess()
{
    {
        std::lock_guard<std::mutex> guard(stateMutex_);
        if (helperProcessValid_) {
            if (WaitForSingleObject(helperProcessInfo_.hProcess, 0) == WAIT_TIMEOUT) {
                hid_handler.log << "Will not launch the alignment helper as it is already running" << std::endl;
                return false;
            }

            CleanupHelperProcessLocked();
        }
    }
    // If an existing ETCalOverlay instance is running, ask it to close first.
    const auto existing = FindProcessIdsByName(kAlignmentProcess);
    for (DWORD pid : existing)
    {
        RequestCloseThenKill(pid, 1500);
        break;
    }

    const std::wstring workingDir = GetExecutableDirectory() + L"\\eyetracking\\AlignmentHelper";
    const std::wstring exePath = workingDir + L"\\" + L"ETCalOverlay.exe";

    // CreateProcessW requires a mutable buffer for lpCommandLine
    std::vector<wchar_t> commandLine;
    commandLine.reserve(exePath.length() + 3);
    commandLine.assign(exePath.begin(), exePath.end());
    commandLine.push_back(L'\0');

    STARTUPINFOW si{};
    si.cb = sizeof(si);
    PROCESS_INFORMATION pi{};

    if (!CreateProcessW(
        nullptr,
        commandLine.data(),
        nullptr,
        nullptr,
        FALSE,
        0,
        nullptr,
        workingDir.c_str(),
        &si,
        &pi))
    {
        hid_handler.log << "Failed to start alignment helper. Win32 error: " << GetLastError() << std::endl;
        return false;
    }

    {
        std::lock_guard<std::mutex> guard(stateMutex_);
        helperProcessInfo_ = pi;
        helperProcessValid_ = true;
    }

    hid_handler.log << "Alignment helper launched successfully" << std::endl;
    return true;
}

bool SteamVRWatch::CloseAlignmentProcess()
{
    std::lock_guard<std::mutex> guard(stateMutex_);

    try
    {
        if (helperProcessValid_ && WaitForSingleObject(helperProcessInfo_.hProcess, 0) == WAIT_TIMEOUT)
        {
            TerminateProcess(helperProcessInfo_.hProcess, 0);
            WaitForSingleObject(helperProcessInfo_.hProcess, 2000);
            hid_handler.log << "Alignment helper closed" << std::endl;
        }

        CleanupHelperProcessLocked();
        return true;
    }
    catch (...)
    {
        hid_handler.log << "Failed to close alignment helper" << std::endl;
        return false;
    }
}

bool SteamVRWatch::IsHelperRunning()
{
    std::lock_guard<std::mutex> guard(stateMutex_);

    if (!helperProcessValid_)
    {
        return false;
    }

    if (WaitForSingleObject(helperProcessInfo_.hProcess, 0) == WAIT_TIMEOUT)
    {
        return true;
    }

    CleanupHelperProcessLocked();
    return false;
}

void SteamVRWatch::CleanupHelperProcessLocked()
{
    if (helperProcessInfo_.hThread)
    {
        CloseHandle(helperProcessInfo_.hThread);
        helperProcessInfo_.hThread = nullptr;
    }

    if (helperProcessInfo_.hProcess)
    {
        CloseHandle(helperProcessInfo_.hProcess);
        helperProcessInfo_.hProcess = nullptr;
    }

    helperProcessInfo_.dwProcessId = 0;
    helperProcessInfo_.dwThreadId = 0;
    helperProcessValid_ = false;
}

bool SteamVRWatch::IsProcessRunningByName(const std::wstring& processNameNoExtension)
{
    return !FindProcessIdsByName(processNameNoExtension).empty();
}

std::wstring SteamVRWatch::GetExecutableDirectory()
{
    wchar_t path[MAX_PATH]{};
    const DWORD length = GetModuleFileNameW(nullptr, path, MAX_PATH);
    if (length == 0 || length == MAX_PATH)
    {
        return L".";
    }

    std::wstring fullPath(path, length);
    const size_t slashPos = fullPath.find_last_of(L"\\/");
    if (slashPos == std::wstring::npos)
    {
        return L".";
    }

    return fullPath.substr(0, slashPos);
}

void SteamVRWatch::UpdateRunningApps()
{
    std::vector<RunningApp> foundApps;
    HANDLE snapshot = CreateToolhelp32Snapshot(
        TH32CS_SNAPPROCESS,
        0);

    if (snapshot == INVALID_HANDLE_VALUE)
        return;

    PROCESSENTRY32W pe{};
    pe.dwSize = sizeof(pe);

    if (Process32FirstW(snapshot, &pe))
    {
        do
        {
            std::wstring title;

            if (ProcessHasVisibleWindow(pe.th32ProcessID, title))
            {
                std::wstring processName = NormalizeProcessName(pe.szExeFile);

                if (processName == L"beyondhid")
                {
                    continue;
                }

                RunningApp app;
                app.processName = processName;
                app.windowTitle = title;

                foundApps.push_back(app);
            }

        } while (Process32NextW(snapshot, &pe));
    }

    CloseHandle(snapshot);
    runningApps = foundApps;
}