#ifndef WIN32_LEAN_AND_MEAN
#define WIN32_LEAN_AND_MEAN
#endif
#ifndef NOMINMAX
#define NOMINMAX
#endif
#include <Windows.h>
#include <shellapi.h>
#include <stdio.h>

#include <Core/Common/CoreUtils.h>
#include <Core/Main.h>
#include <Core/Memory/MemoryArena.h>
#include <Core/JulietInit.h>
#include <Core/HAL/Filesystem/Filesystem.h>
#include <Core/Thread/Mutex.h>
#include <thread>
#include "Database.h"
#include "MarkdownGenerator.h"

#define ROMEO_TRAY_ICON_ID 1
#define ROMEO_WM_TRAYICON (WM_USER + 1)

static Romeo::Database G_Db;
static Juliet::Arena* G_ScratchArena = nullptr;
static Juliet::String G_WorkspaceRoot;
static Juliet::String G_DocsDir;
static Juliet::String G_DbFilePath;

static Juliet::String GetWorkspaceRoot(Juliet::Arena* arena)
{
    Assert(arena != nullptr);
    Juliet::String base = Juliet::GetBasePath();
    
    // Copy to arena so we can safely modify it
    char* buf = Juliet::ArenaPushArray<char>(arena, base.Size + 1 JULIET_DEBUG_PARAM("WorkspaceRootPath"));
    memcpy(buf, base.Data, base.Size);
    buf[base.Size] = '\0';

    // Normalize slashes
    for (size_t i = 0; i < base.Size; ++i)
    {
        if (buf[i] == '/')
        {
            buf[i] = '\\';
        }
    }

    // Search for \bin to find workspace root
    char* binPtr = strstr(buf, "\\bin");
    if (binPtr != nullptr)
    {
        *binPtr = '\0';
        return { buf, static_cast<size_t>(binPtr - buf) };
    }

    size_t count = 0;
    for (size_t i = base.Size; i > 0; --i)
    {
        size_t idx = i - 1;
        if (buf[idx] == '\\')
        {
            count++;
            if (count == 2)
            {
                buf[idx] = '\0';
                return { buf, idx };
            }
        }
    }
    return { buf, base.Size };
}

static HANDLE G_ExitEvent = INVALID_HANDLE_VALUE;
static std::thread G_MonitorThread;
static bool G_MonitorRunning = false;
static HANDLE G_OllamaProcessHandle = INVALID_HANDLE_VALUE;
static HANDLE G_OllamaJobHandle = nullptr;

static void MonitorThreadProc()
{
    int wLen = MultiByteToWideChar(CP_UTF8, 0, CStr(G_WorkspaceRoot), -1, nullptr, 0);
    if (wLen <= 0)
    {
        return;
    }
    
    wchar_t* wRoot = new wchar_t[static_cast<size_t>(wLen + 1)];
    MultiByteToWideChar(CP_UTF8, 0, CStr(G_WorkspaceRoot), -1, wRoot, wLen);
    
    HANDLE hDir = CreateFileW(
        wRoot,
        FILE_LIST_DIRECTORY,
        FILE_SHARE_READ | FILE_SHARE_WRITE | FILE_SHARE_DELETE,
        nullptr,
        OPEN_EXISTING,
        FILE_FLAG_BACKUP_SEMANTICS | FILE_FLAG_OVERLAPPED,
        nullptr
    );
    
    delete[] wRoot;
    
    if (hDir == INVALID_HANDLE_VALUE)
    {
        return;
    }
    
    HANDLE hEvent = CreateEventW(nullptr, TRUE, FALSE, nullptr);
    if (hEvent == nullptr)
    {
        CloseHandle(hDir);
        return;
    }
    
    OVERLAPPED overlapped = {};
    overlapped.hEvent = hEvent;
    
    alignas(FILE_NOTIFY_INFORMATION) BYTE buffer[64 * 1024];
    bool pendingChanges = false;
    
    BOOL readRes = ReadDirectoryChangesW(
        hDir,
        buffer,
        sizeof(buffer),
        TRUE,
        FILE_NOTIFY_CHANGE_LAST_WRITE | FILE_NOTIFY_CHANGE_FILE_NAME,
        nullptr,
        &overlapped,
        nullptr
    );
    
    if (readRes == 0)
    {
        CloseHandle(hEvent);
        CloseHandle(hDir);
        return;
    }
    
    HANDLE waitHandles[2] = { G_ExitEvent, hEvent };
    
    while (G_MonitorRunning)
    {
        DWORD timeout = pendingChanges ? 500 : INFINITE;
        DWORD waitRes = WaitForMultipleObjects(2, waitHandles, FALSE, timeout);
        
        if (waitRes == WAIT_OBJECT_0)
        {
            break;
        }
        else if (waitRes == WAIT_OBJECT_0 + 1)
        {
            DWORD bytesTransferred = 0;
            if (GetOverlappedResult(hDir, &overlapped, &bytesTransferred, FALSE) != 0)
            {
                FILE_NOTIFY_INFORMATION* info = reinterpret_cast<FILE_NOTIFY_INFORMATION*>(buffer);
                while (info != nullptr)
                {
                    size_t wCharCount = info->FileNameLength / sizeof(WCHAR);
                    if (wCharCount > 2)
                    {
                        const WCHAR* fName = info->FileName;
                        if ((fName[wCharCount - 1] == L'h' && fName[wCharCount - 2] == L'.') ||
                            (wCharCount > 4 && fName[wCharCount - 1] == L'p' && fName[wCharCount - 2] == L'p' && fName[wCharCount - 3] == L'c' && fName[wCharCount - 4] == L'.'))
                        {
                            bool ignore = false;
                            for (size_t i = 0; i < wCharCount; ++i)
                            {
                                if (i + 12 <= wCharCount && wcsncmp(fName + i, L"Intermediate", 12) == 0) { ignore = true; }
                                if (i + 3 <= wCharCount && wcsncmp(fName + i, L"bin", 3) == 0) { ignore = true; }
                                if (i + 8 <= wCharCount && wcsncmp(fName + i, L"External", 8) == 0) { ignore = true; }
                                if (i + 4 <= wCharCount && wcsncmp(fName + i, L".git", 4) == 0) { ignore = true; }
                                if (ignore) { break; }
                            }
                            
                            if (!ignore)
                            {
                                pendingChanges = true;
                            }
                        }
                    }
                    
                    if (info->NextEntryOffset == 0)
                    {
                        break;
                    }
                    info = reinterpret_cast<FILE_NOTIFY_INFORMATION*>(reinterpret_cast<BYTE*>(info) + info->NextEntryOffset);
                }
            }
            
            ResetEvent(hEvent);
            ReadDirectoryChangesW(
                hDir,
                buffer,
                sizeof(buffer),
                TRUE,
                FILE_NOTIFY_CHANGE_LAST_WRITE | FILE_NOTIFY_CHANGE_FILE_NAME,
                nullptr,
                &overlapped,
                nullptr
            );
        }
        else if (waitRes == WAIT_TIMEOUT)
        {
            printf("Romeo Monitor: Change detected, updating documentation...\n");
            OutputDebugStringA("Romeo Monitor: Change detected, updating...\n");
            
            {
                Juliet::LockGuard lock(G_Db.Mutex);
                Juliet::ArenaClear(G_ScratchArena);
                
                Romeo::DB_ScanWorkspace(G_Db, G_ScratchArena, G_WorkspaceRoot);
                
                bool anyChanged = false;
                for (const Romeo::FileEntry& entry : G_Db.Entries)
                {
                    if (entry.LastModifiedHeader > entry.LastMarkdownCreatedTime ||
                        entry.LastModifiedCpp > entry.LastMarkdownCreatedTime)
                    {
                        anyChanged = true;
                        break;
                    }
                }
                
                if (anyChanged)
                {
                    Romeo::MD_GenerateDocumentation(G_Db, G_WorkspaceRoot, G_DocsDir);
                    Romeo::DB_Save(G_Db, G_DbFilePath);
                }
            }
            
            pendingChanges = false;
        }
    }
    
    CancelIo(hDir);
    CloseHandle(hEvent);
    CloseHandle(hDir);
}

LRESULT CALLBACK RomeoWindowProc(HWND hwnd, UINT uMsg, WPARAM wParam, LPARAM lParam)
{
    switch (uMsg)
    {
    case ROMEO_WM_TRAYICON:
        if (LOWORD(lParam) == WM_RBUTTONUP)
        {
            POINT pt;
            GetCursorPos(&pt);
            HMENU hMenu = CreatePopupMenu();
            if (hMenu)
            {
                InsertMenuW(hMenu, 0, MF_BYPOSITION | MF_STRING, 1001, L"Exit");
                SetForegroundWindow(hwnd);
                TrackPopupMenu(hMenu, TPM_BOTTOMALIGN | TPM_LEFTALIGN, pt.x, pt.y, 0, hwnd, nullptr);
                DestroyMenu(hMenu);
            }
        }
        else if (LOWORD(lParam) == WM_LBUTTONUP)
        {
            OutputDebugStringA("Tray icon left-clicked!\n");
        }
        return 0;

    case WM_COMMAND:
        if (LOWORD(wParam) == 1001)
        {
            DestroyWindow(hwnd);
        }
        return 0;

    case WM_DESTROY:
        PostQuitMessage(0);
        return 0;
    }
    return DefWindowProc(hwnd, uMsg, wParam, lParam);
}

int JulietMain(int argc, wchar_t** argv)
{
    printf("JulietMain Entry!\n");

    // Single instance mutex check
    HANDLE hMutex = CreateMutexW(nullptr, TRUE, L"RomeoSingleInstanceMutex");
    if (hMutex == nullptr || GetLastError() == ERROR_ALREADY_EXISTS)
    {
        printf("Romeo is already running. Exiting.\n");
        if (hMutex != nullptr)
        {
            CloseHandle(hMutex);
        }
        return 0;
    }

    // Launch Ollama serve subprocess
    {
        // Try to find ollama.exe - check common install locations
        char ollamaPath[512] = {};
        bool foundOllama = false;

        // 1. Try PATH first via SearchPathA
        if (SearchPathA(nullptr, "ollama.exe", nullptr, sizeof(ollamaPath), ollamaPath, nullptr) > 0)
        {
            foundOllama = true;
        }

        // 2. Check %LOCALAPPDATA%\Programs\Ollama\ollama.exe
        if (!foundOllama)
        {
            char localAppData[256] = {};
            DWORD envLen = GetEnvironmentVariableA("LOCALAPPDATA", localAppData, sizeof(localAppData));
            if (envLen > 0 && envLen < sizeof(localAppData))
            {
                snprintf(ollamaPath, sizeof(ollamaPath), "%s\\Programs\\Ollama\\ollama.exe", localAppData);
                DWORD attribs = GetFileAttributesA(ollamaPath);
                if (attribs != INVALID_FILE_ATTRIBUTES)
                {
                    foundOllama = true;
                }
            }
        }

        // 3. Check %USERPROFILE%\AppData\Local\Programs\Ollama\ollama.exe
        if (!foundOllama)
        {
            char userProfile[256] = {};
            DWORD envLen = GetEnvironmentVariableA("USERPROFILE", userProfile, sizeof(userProfile));
            if (envLen > 0 && envLen < sizeof(userProfile))
            {
                snprintf(ollamaPath, sizeof(ollamaPath), "%s\\AppData\\Local\\Programs\\Ollama\\ollama.exe", userProfile);
                DWORD attribs = GetFileAttributesA(ollamaPath);
                if (attribs != INVALID_FILE_ATTRIBUTES)
                {
                    foundOllama = true;
                }
            }
        }

        if (foundOllama)
        {
            printf("Romeo: Found Ollama at: %s\n", ollamaPath);

            // Create a Job Object so child processes auto-terminate when Romeo exits (even on crash)
            G_OllamaJobHandle = CreateJobObjectA(nullptr, nullptr);
            if (G_OllamaJobHandle != nullptr)
            {
                JOBOBJECT_EXTENDED_LIMIT_INFORMATION jobInfo = {};
                jobInfo.BasicLimitInformation.LimitFlags = JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE;
                SetInformationJobObject(G_OllamaJobHandle, JobObjectExtendedLimitInformation,
                                        &jobInfo, sizeof(jobInfo));
            }

            // Build a command line for cmd.exe /k so the console window remains open
            char cmdLine[600] = {};
            snprintf(cmdLine, sizeof(cmdLine), "cmd.exe /k \"\"%s\" serve\"", ollamaPath);

            STARTUPINFOA si = {};
            si.cb = sizeof(si);
            si.dwFlags = STARTF_USESHOWWINDOW;
            si.wShowWindow = SW_SHOW;
            PROCESS_INFORMATION pi = {};
            BOOL createRes = CreateProcessA(
                nullptr,
                cmdLine,
                nullptr,
                nullptr,
                FALSE,
                CREATE_NEW_CONSOLE | CREATE_SUSPENDED,
                nullptr,
                nullptr,
                &si,
                &pi
            );
            if (createRes != 0)
            {
                // Assign to job object before resuming so it cannot escape
                if (G_OllamaJobHandle != nullptr)
                {
                    AssignProcessToJobObject(G_OllamaJobHandle, pi.hProcess);
                }
                ResumeThread(pi.hThread);

                // Wait briefly to check if the process crashed immediately
                DWORD waitResult = WaitForSingleObject(pi.hProcess, 2000);
                if (waitResult == WAIT_TIMEOUT)
                {
                    // Process is still running after 2 seconds - good
                    G_OllamaProcessHandle = pi.hProcess;
                    CloseHandle(pi.hThread);
                    printf("Romeo: Launched local Ollama serve subprocess (PID: %lu).\n", pi.dwProcessId);
                }
                else
                {
                    // Process exited within 2 seconds - likely crashed or port in use
                    DWORD exitCode = 0;
                    GetExitCodeProcess(pi.hProcess, &exitCode);
                    printf("Romeo: WARNING - Ollama serve process exited immediately (exit code: %lu).\n", exitCode);
                    printf("Romeo: This may mean Ollama is already running or the port 11434 is in use.\n");
                    CloseHandle(pi.hProcess);
                    CloseHandle(pi.hThread);
                }
            }
            else
            {
                DWORD err = GetLastError();
                printf("Romeo: ERROR - Failed to launch Ollama serve (Win32 error: %lu).\n", err);
            }
        }
        else
        {
            printf("Romeo: WARNING - Could not find ollama.exe. Searched PATH and common install locations.\n");
            printf("Romeo: Please ensure Ollama is installed (https://ollama.com) and available in your PATH.\n");
            printf("Romeo: Continuing without auto-launching Ollama serve...\n");
        }
    }

    Juliet::ArenaParams params{};
    Juliet::Arena* mainArena = Juliet::ArenaAllocate(params JULIET_DEBUG_PARAM("RomeoMainArena"));
    Assert(mainArena != nullptr);

    G_ScratchArena = Juliet::ArenaAllocate(params JULIET_DEBUG_PARAM("RomeoScratchArena"));
    Assert(G_ScratchArena != nullptr);

    printf("Arena Allocated! Pointer: %p\n", static_cast<void*>(mainArena));

    G_WorkspaceRoot = GetWorkspaceRoot(mainArena);

    char docsDirBuf[512];
    snprintf(docsDirBuf, sizeof(docsDirBuf), "%s\\Romeo\\docs", CStr(G_WorkspaceRoot));
    G_DocsDir = Juliet::StringCopy(mainArena, Juliet::WrapString(docsDirBuf));

    char dbFilePathBuf[512];
    snprintf(dbFilePathBuf, sizeof(dbFilePathBuf), "%s\\Romeo\\docs\\romeo.db", CStr(G_WorkspaceRoot));
    G_DbFilePath = Juliet::StringCopy(mainArena, Juliet::WrapString(dbFilePathBuf));

    Romeo::DB_Init(G_Db, mainArena);

    printf("DB Init!\n");

    bool forceRegenerate = false;
    if (argc > 1 && (wcscmp(argv[1], L"-force") == 0 || wcscmp(argv[1], L"-f") == 0))
    {
        forceRegenerate = true;
    }

    {
        Juliet::LockGuard lock(G_Db.Mutex);
        if (!forceRegenerate)
        {
            Romeo::DB_Load(G_Db, G_DbFilePath);
        }
        else
        {
            printf("Romeo: Force flag detected, clearing cache and rebuilding all documentation...\n");
        }

        Romeo::DB_ScanWorkspace(G_Db, G_ScratchArena, G_WorkspaceRoot);

        if (forceRegenerate)
        {
            for (Romeo::FileEntry& entry : G_Db.Entries)
            {
                entry.Description = {};
                entry.LastMarkdownCreatedTime = 0;
                entry.LastAIGeneratedTime = 0;
            }
        }

        Romeo::MD_GenerateDocumentation(G_Db, G_WorkspaceRoot, G_DocsDir);
        Romeo::DB_Save(G_Db, G_DbFilePath);
    }

    char debugMsg[256];
    snprintf(debugMsg, sizeof(debugMsg), "Romeo scanned workspace. Found %llu C++ files.\n", static_cast<unsigned long long>(G_Db.Entries.Size()));
    printf("%s", debugMsg);

    // Register hidden window class
    const wchar_t CLASS_NAME[] = L"RomeoHiddenWindowClass";
    WNDCLASSW wc = {};
    HINSTANCE hInstance = GetModuleHandleW(nullptr);
    wc.lpfnWndProc = RomeoWindowProc;
    wc.hInstance = hInstance;
    wc.lpszClassName = CLASS_NAME;

    RegisterClassW(&wc);

    // Create hidden window to receive messages
    HWND hwnd = CreateWindowExW(
        0,
        CLASS_NAME,
        L"Romeo Tool",
        0, 
        CW_USEDEFAULT, CW_USEDEFAULT, CW_USEDEFAULT, CW_USEDEFAULT,
        nullptr,
        nullptr,
        hInstance,
        nullptr
    );

    if (hwnd == nullptr)
    {
        CloseHandle(hMutex);
        return 0;
    }

    G_ExitEvent = CreateEventW(nullptr, TRUE, FALSE, nullptr);
    if (G_ExitEvent != INVALID_HANDLE_VALUE)
    {
        G_MonitorRunning = true;
        G_MonitorThread = std::thread(MonitorThreadProc);
    }

    // Add tray icon
    NOTIFYICONDATAW nid = {};
    nid.cbSize = sizeof(NOTIFYICONDATAW);
    nid.hWnd = hwnd;
    nid.uID = ROMEO_TRAY_ICON_ID;
    nid.uFlags = NIF_ICON | NIF_MESSAGE | NIF_TIP;
    nid.uCallbackMessage = ROMEO_WM_TRAYICON;
    nid.hIcon = LoadIcon(nullptr, IDI_APPLICATION);
    wcscpy_s(nid.szTip, L"Romeo Documentation Tool");

    Shell_NotifyIconW(NIM_ADD, &nid);

    // Message loop
    MSG msg = {};
    while (GetMessage(&msg, nullptr, 0, 0))
    {
        TranslateMessage(&msg);
        DispatchMessage(&msg);
    }

    // Cleanup
    Shell_NotifyIconW(NIM_DELETE, &nid);

    G_MonitorRunning = false;
    if (G_ExitEvent != INVALID_HANDLE_VALUE)
    {
        SetEvent(G_ExitEvent);
    }
    if (G_MonitorThread.joinable())
    {
        G_MonitorThread.join();
    }
    if (G_ExitEvent != INVALID_HANDLE_VALUE)
    {
        CloseHandle(G_ExitEvent);
        G_ExitEvent = INVALID_HANDLE_VALUE;
    }

    if (G_OllamaProcessHandle != INVALID_HANDLE_VALUE)
    {
        TerminateProcess(G_OllamaProcessHandle, 0);
        CloseHandle(G_OllamaProcessHandle);
        G_OllamaProcessHandle = INVALID_HANDLE_VALUE;
    }
    if (G_OllamaJobHandle != nullptr)
    {
        CloseHandle(G_OllamaJobHandle);
        G_OllamaJobHandle = nullptr;
    }

    Juliet::ArenaRelease(mainArena);
    Juliet::ArenaRelease(G_ScratchArena);

    CloseHandle(hMutex);

    return 0;
}
