#include "AiGenerator.h"
#include "HttpClient.h"
#include "JsonParser.h"
#include <Core/Common/CoreUtils.h>
#include <Core/HAL/IO/IOStream.h>
#include <Core/Memory/MemoryArena.h>
#include <stdio.h>
#include <string.h>
#include <Windows.h>

namespace Romeo
{
    static Juliet::String ReadFileContent(Juliet::Arena* arena, const Juliet::String& path)
    {
        Assert(arena != nullptr);
        if (path.Size == 0)
        {
            return {};
        }
        
        char pathBuf[512];
        size_t copyLen = path.Size < 511 ? path.Size : 511;
        memcpy(pathBuf, path.Data, copyLen);
        pathBuf[copyLen] = '\0';
        
        FILE* f = nullptr;
        if (fopen_s(&f, pathBuf, "rb") != 0 || !f)
        {
            return {};
        }

        fseek(f, 0, SEEK_END);
        long size = ftell(f);
        fseek(f, 0, SEEK_SET);

        if (size <= 0)
        {
            fclose(f);
            return {};
        }
        
        // Limit to 20KB to keep prompt token size reasonable
        if (size > 20480)
        {
            size = 20480;
        }
        
        char* buf = Juliet::ArenaPushArray<char>(arena, static_cast<size_t>(size + 1) JULIET_DEBUG_PARAM("FileReadContent"));
        size_t readBytes = fread(buf, 1, static_cast<size_t>(size), f);
        buf[readBytes] = '\0';
        
        fclose(f);
        return { buf, readBytes };
    }

    static Juliet::String GetEnvironmentVar(Juliet::Arena* arena, const char* varName)
    {
        Assert(arena != nullptr);
        Assert(varName != nullptr);
        
        DWORD size = GetEnvironmentVariableA(varName, nullptr, 0);
        if (size == 0)
        {
            return {};
        }
        
        char* buf = Juliet::ArenaPushArray<char>(arena, size JULIET_DEBUG_PARAM("EnvVar"));
        GetEnvironmentVariableA(varName, buf, size);
        return { buf, static_cast<size_t>(size - 1) };
    }

    static bool G_OllamaReady = false;

    static bool PostWithRetry(
        Juliet::Arena* arena,
        const Juliet::String& host,
        int port,
        const Juliet::String& path,
        bool isHttps,
        const Juliet::String& requestBody,
        Juliet::String& outResponse
    )
    {
        if (isHttps)
        {
            return HttpClient::Post(arena, host, port, path, isHttps, requestBody, outResponse);
        }

        if (G_OllamaReady)
        {
            return HttpClient::Post(arena, host, port, path, isHttps, requestBody, outResponse);
        }

        int retries = 15;
        while (retries > 0)
        {
            if (HttpClient::Post(arena, host, port, path, isHttps, requestBody, outResponse))
            {
                G_OllamaReady = true;
                return true;
            }
            retries--;
            if (retries > 0)
            {
                printf("Romeo AI: Local Ollama serve not ready yet. Retrying in 1s (%d retries remaining)...\n", retries);
                Sleep(1000);
            }
        }
        return false;
    }

    static bool PullOllamaModelIfNeeded(Juliet::Arena* arena, const Juliet::String& host, int port, const Juliet::String& modelName)
    {
        size_t bodyCap = modelName.Size + 64;
        char* showBodyBuf = Juliet::ArenaPushArray<char>(arena, bodyCap JULIET_DEBUG_PARAM("ShowBody"));
        int showLen = snprintf(showBodyBuf, bodyCap, "{\"name\":\"%.*s\"}", static_cast<int>(modelName.Size), CStr(modelName));
        Juliet::String showBody = { showBodyBuf, static_cast<size_t>(showLen > 0 ? showLen : 0) };

        Juliet::String showResponse = {};
        bool showSuccess = PostWithRetry(arena, host, port, ConstString("/api/show"), false, showBody, showResponse);

        bool needsPull = false;
        if (!showSuccess || showResponse.Size == 0)
        {
            needsPull = true;
        }
        else
        {
            JsonValue root = {};
            if (Json_Parse(arena, showResponse, root))
            {
                Juliet::String errorMsg = {};
                if (Json_Query(root, "error", errorMsg))
                {
                    needsPull = true;
                }
            }
            else
            {
                needsPull = true;
            }
        }

        if (needsPull)
        {
            printf("Romeo AI: Ollama model '%.*s' not found. Pulling model automatically (this may take a few minutes)...\n",
                   static_cast<int>(modelName.Size), CStr(modelName));
            OutputDebugStringA("Romeo AI: Model not found. Pulling model...\n");

            char* pullBodyBuf = Juliet::ArenaPushArray<char>(arena, bodyCap + 32 JULIET_DEBUG_PARAM("PullBody"));
            int pullLen = snprintf(pullBodyBuf, bodyCap + 32, "{\"name\":\"%.*s\",\"stream\":false}", static_cast<int>(modelName.Size), CStr(modelName));
            Juliet::String pullBody = { pullBodyBuf, static_cast<size_t>(pullLen > 0 ? pullLen : 0) };

            Juliet::String pullResponse = {};
            bool pullSuccess = PostWithRetry(arena, host, port, ConstString("/api/pull"), false, pullBody, pullResponse);
            if (pullSuccess)
            {
                printf("Romeo AI: Model '%.*s' pulled successfully!\n", static_cast<int>(modelName.Size), CStr(modelName));
                OutputDebugStringA("Romeo AI: Model pulled successfully!\n");
                return true;
            }
            else
            {
                printf("Romeo AI: Failed to pull model '%.*s' from Ollama API.\n", static_cast<int>(modelName.Size), CStr(modelName));
                OutputDebugStringA("Romeo AI: Failed to pull model.\n");
                return false;
            }
        }

        return true;
    }

    static Juliet::String CleanAndFormatDescription(Juliet::Arena* arena, const Juliet::String& rawResponse, bool enableCoT)
    {
        Assert(arena != nullptr);
        if (rawResponse.Size == 0)
        {
            return {};
        }

        const char* pStr = CStr(rawResponse);
        size_t len = rawResponse.Size;

        const char* thinkStart = strstr(pStr, "<think>");
        const char* thinkEnd = thinkStart ? strstr(thinkStart, "</think>") : nullptr;

        Juliet::String thinkContent = {};
        Juliet::String mainContent = {};

        if (thinkStart && thinkEnd && thinkEnd > thinkStart)
        {
            const char* thinkTextStart = thinkStart + 7;
            size_t thinkLen = static_cast<size_t>(thinkEnd - thinkTextStart);
            thinkContent = { const_cast<char*>(thinkTextStart), thinkLen };

            const char* afterThink = thinkEnd + 8;
            size_t remainingLen = len - static_cast<size_t>(afterThink - pStr);
            mainContent = { const_cast<char*>(afterThink), remainingLen };
        }
        else
        {
            mainContent = rawResponse;
        }

        // Clean mainContent (remove titles, markdown headers, leading quotes/fences)
        const char* mPtr = CStr(mainContent);
        size_t mLen = mainContent.Size;

        while (mLen > 0 && (*mPtr == ' ' || *mPtr == '\t' || *mPtr == '\r' || *mPtr == '\n' || *mPtr == '#' || *mPtr == '`'))
        {
            mPtr++;
            mLen--;
        }

        while (mLen > 0 && (mPtr[mLen - 1] == ' ' || mPtr[mLen - 1] == '\t' || mPtr[mLen - 1] == '\r' || mPtr[mLen - 1] == '\n' || mPtr[mLen - 1] == '`'))
        {
            mLen--;
        }

        if (enableCoT && thinkContent.Size > 0)
        {
            size_t resultCap = thinkContent.Size + mLen + 256;
            char* resultBuf = Juliet::ArenaPushArray<char>(arena, resultCap JULIET_DEBUG_PARAM("CoTFormattedDesc"));
            int written = snprintf(resultBuf, resultCap,
                "<details>\n<summary>Chain of Thought</summary>\n\n%.*s\n\n</details>\n\n%.*s",
                static_cast<int>(thinkContent.Size), CStr(thinkContent),
                static_cast<int>(mLen), mPtr
            );
            return { resultBuf, static_cast<size_t>(written > 0 ? written : 0) };
        }

        char* resultBuf = Juliet::ArenaPushArray<char>(arena, mLen + 1 JULIET_DEBUG_PARAM("CleanedDesc"));
        memcpy(resultBuf, mPtr, mLen);
        resultBuf[mLen] = '\0';
        return { resultBuf, mLen };
    }

    bool AI_GenerateDescription(
        Juliet::Arena* arena,
        FileEntry& entry
    )
    {
        Assert(arena != nullptr);

        Juliet::String headerContent = ReadFileContent(arena, entry.HeaderPath);
        Juliet::String cppContent = ReadFileContent(arena, entry.CppPath);

        if (headerContent.Size == 0 && cppContent.Size == 0)
        {
            entry.Description = Juliet::StringCopy(arena, ConstString("*File is empty.*"));
            return true;
        }

        Juliet::String cotEnv = GetEnvironmentVar(arena, "ROMEO_ENABLE_COT");
        bool enableCoT = (cotEnv.Size > 0 && (cotEnv.Data[0] == '1' || cotEnv.Data[0] == 't' || cotEnv.Data[0] == 'T' || cotEnv.Data[0] == 'y' || cotEnv.Data[0] == 'Y'));

        const char* sysPromptNoCoT = "You are a technical documentation assistant. Provide a single, concise, unified paragraph overview (50-80 words) describing the main purpose and design of this C++ component.\n"
            "Rules:\n"
            "1. Do NOT split your output by header file or source file.\n"
            "2. Do NOT include titles, markdown headers (e.g. # or ##), bullet lists, or code blocks.\n"
            "3. Do NOT include <think> tags or reasoning text.\n"
            "4. Return ONLY the description paragraph.";

        const char* sysPromptCoT = "You are a technical documentation assistant. Provide a single, concise, unified paragraph overview (50-80 words) describing the main purpose and design of this C++ component.\n"
            "Rules:\n"
            "1. You may include your reasoning inside <think>...</think> tags.\n"
            "2. Following the reasoning, output ONLY the final description paragraph.\n"
            "3. Do NOT split your final output by header file or source file.";

        const char* sysPrompt = enableCoT ? sysPromptCoT : sysPromptNoCoT;
        
        size_t promptCap = strlen(sysPrompt) + entry.HeaderPath.Size + headerContent.Size + entry.CppPath.Size + cppContent.Size + entry.Description.Size + 512;
        char* promptBuf = Juliet::ArenaPushArray<char>(arena, promptCap JULIET_DEBUG_PARAM("PromptBuf"));
        
        int written = 0;
        if (entry.Description.Size > 0)
        {
            written = snprintf(promptBuf, promptCap,
                "%s\n\n"
                "Current description:\n%.*s\n\n"
                "Header file: %.*s\n"
                "%.*s\n\n"
                "Source file: %.*s\n"
                "%.*s\n",
                sysPrompt,
                static_cast<int>(entry.Description.Size), CStr(entry.Description),
                static_cast<int>(entry.HeaderPath.Size), CStr(entry.HeaderPath),
                static_cast<int>(headerContent.Size), CStr(headerContent),
                static_cast<int>(entry.CppPath.Size), CStr(entry.CppPath),
                static_cast<int>(cppContent.Size), CStr(cppContent)
            );
        }
        else
        {
            written = snprintf(promptBuf, promptCap,
                "%s\n\n"
                "Header file: %.*s\n"
                "%.*s\n\n"
                "Source file: %.*s\n"
                "%.*s\n",
                sysPrompt,
                static_cast<int>(entry.HeaderPath.Size), CStr(entry.HeaderPath),
                static_cast<int>(headerContent.Size), CStr(headerContent),
                static_cast<int>(entry.CppPath.Size), CStr(entry.CppPath),
                static_cast<int>(cppContent.Size), CStr(cppContent)
            );
        }

        Juliet::String prompt = { promptBuf, static_cast<size_t>(written > 0 ? written : 0) };
        Juliet::String escapedPrompt = Json_EscapeString(arena, prompt);

        Juliet::String host = ConstString("127.0.0.1");
        int port = 11434;
        Juliet::String modelName = GetEnvironmentVar(arena, "ROMEO_LLM_MODEL");
        if (modelName.Size == 0)
        {
            modelName = ConstString("qwen2.5-coder:1.5b");
        }

        if (!PullOllamaModelIfNeeded(arena, host, port, modelName))
        {
            return false;
        }

        printf("Romeo AI: Querying local Ollama (%.*s) [CoT: %s] for '%.*s'...\n",
               static_cast<int>(modelName.Size), CStr(modelName),
               enableCoT ? "ON" : "OFF",
               static_cast<int>(entry.HeaderPath.Size > 0 ? entry.HeaderPath.Size : entry.CppPath.Size),
               CStr(entry.HeaderPath.Size > 0 ? entry.HeaderPath : entry.CppPath));

        int maxPredictTokens = enableCoT ? 600 : 300;
        size_t bodyCap = modelName.Size + escapedPrompt.Size + 512;
        char* bodyBuf = Juliet::ArenaPushArray<char>(arena, bodyCap JULIET_DEBUG_PARAM("OllamaBody"));
        int bodyLen = snprintf(bodyBuf, bodyCap,
            "{\"model\":\"%.*s\",\"prompt\":\"%.*s\",\"stream\":false,\"options\":{\"temperature\":0.2,\"num_predict\":%d}}",
            static_cast<int>(modelName.Size), CStr(modelName),
            static_cast<int>(escapedPrompt.Size), CStr(escapedPrompt),
            maxPredictTokens
        );
        Juliet::String body = { bodyBuf, static_cast<size_t>(bodyLen > 0 ? bodyLen : 0) };

        Juliet::String response = {};
        bool success = PostWithRetry(
            arena,
            host,
            port,
            ConstString("/api/generate"),
            false,
            body,
            response
        );

        if (success && response.Size > 0)
        {
            JsonValue root = {};
            if (Json_Parse(arena, response, root))
            {
                Juliet::String desc = {};
                if (Json_Query(root, "response", desc))
                {
                    entry.Description = CleanAndFormatDescription(arena, desc, enableCoT);
                    return true;
                }
            }
        }

        printf("Romeo AI: Failed to communicate with local Ollama.\n");
        return false;
    }
}
