#include "LLMService.hpp"
#include "../shared/llm/LLMProviderFactory.hpp"

#include <spdlog/sinks/stdout_color_sinks.h>
#include <fstream>

namespace celuna {

LLMService::LLMService() {
    m_logger = spdlog::get("LLMService");
    if (!m_logger) {
        m_logger = spdlog::stdout_color_mt("LLMService");
    }
}

LLMService::~LLMService() {
    shutdown();
}

bool LLMService::initialize(grove::IIO* io) {
    m_io = io;

    if (m_io) {
        grove::SubscriptionConfig config;
        m_io->subscribe("llm:request", config);
    }

    // Start worker thread
    m_running = true;
    m_workerThread = std::thread(&LLMService::workerLoop, this);

    m_logger->info("LLMService initialized");
    return true;
}

bool LLMService::loadConfig(const std::string& configPath) {
    try {
        std::ifstream file(configPath);
        if (!file.is_open()) {
            m_logger->warn("Config file not found: {}", configPath);
            return false;
        }

        nlohmann::json config;
        file >> config;

        m_provider = LLMProviderFactory::create(config);
        if (!m_provider) {
            m_logger->error("Failed to create LLM provider");
            return false;
        }

        m_providerName = config.value("provider", "claude");
        m_maxIterations = config.value("max_iterations", 10);
        m_defaultSystemPrompt = config.value("system_prompt",
            "Tu es AISSIA, un assistant personnel intelligent. "
            "Tu peux utiliser des tools pour accomplir des taches: "
            "gerer le planning, verifier le focus, sauvegarder des notes, "
            "lire des fichiers, faire des recherches web, etc.");

        m_logger->info("LLM provider loaded: {} ({})", m_providerName, m_provider->getModel());

        // Initialize tools after provider is ready
        initializeTools();

        return true;

    } catch (const std::exception& e) {
        m_logger->error("Failed to load config: {}", e.what());
        return false;
    }
}

void LLMService::initializeTools() {
    m_logger->info("Initializing tools...");

    // 1. Internal tools (via GroveEngine IIO)
    if (m_io) {
        m_internalTools = std::make_unique<InternalTools>(m_io);
        for (const auto& tool : m_internalTools->getTools()) {
            m_toolRegistry.registerTool(tool);
        }
        m_logger->info("Registered {} internal tools", m_internalTools->size());
    }

    // 2. FileSystem tools (direct C++ execution)
    for (const auto& toolDef : tools::FileSystemTools::getToolDefinitions()) {
        std::string toolName = toolDef["name"].get<std::string>();
        m_toolRegistry.registerTool(
            toolName,
            toolDef["description"].get<std::string>(),
            toolDef["input_schema"],
            [toolName](const nlohmann::json& input) -> nlohmann::json {
                return tools::FileSystemTools::execute(toolName, input);
            }
        );
    }
    m_logger->info("Registered {} filesystem tools", tools::FileSystemTools::getToolDefinitions().size());

    // 3. MCP tools (via external servers)
    m_mcpClient = std::make_unique<mcp::MCPClient>();
    if (loadMCPConfig("config/mcp.json")) {
        int connected = m_mcpClient->connectAll();
        if (connected > 0) {
            for (const auto& tool : m_mcpClient->listAllTools()) {
                // Convert MCP tool to our ToolDefinition format
                m_toolRegistry.registerTool(
                    tool.name,
                    tool.description,
                    tool.inputSchema,
                    [this, toolName = tool.name](const nlohmann::json& input) -> nlohmann::json {
                        auto result = m_mcpClient->callTool(toolName, input);
                        // Convert MCP result to simple JSON
                        if (result.isError) {
                            return {{"error", true}, {"content", result.content}};
                        }
                        // Extract text content
                        std::string text;
                        for (const auto& content : result.content) {
                            if (content.contains("text")) {
                                text += content["text"].get<std::string>();
                            }
                        }
                        return {{"content", text}};
                    }
                );
            }
            m_logger->info("Registered {} MCP tools from {} servers",
                           m_mcpClient->toolCount(), connected);
        }
    }

    m_logger->info("Total tools available: {}", m_toolRegistry.size());
}

bool LLMService::loadMCPConfig(const std::string& configPath) {
    return m_mcpClient->loadConfig(configPath);
}

void LLMService::registerTool(const std::string& name, const std::string& description,
                              const nlohmann::json& schema,
                              std::function<nlohmann::json(const nlohmann::json&)> handler) {
    m_toolRegistry.registerTool(name, description, schema, handler);
    m_logger->debug("Tool registered: {}", name);
}

void LLMService::process() {
    processIncomingMessages();
    publishResponses();
}

void LLMService::processIncomingMessages() {
    if (!m_io) return;

    while (m_io->hasMessages() > 0) {
        auto msg = m_io->pullMessage();

        if (msg.topic == "llm:request" && msg.data) {
            Request req;
            req.query = msg.data->getString("query", "");
            req.systemPrompt = msg.data->getString("systemPrompt", m_defaultSystemPrompt);
            req.conversationId = msg.data->getString("conversationId", "default");
            req.maxIterations = msg.data->getInt("maxIterations", m_maxIterations);

            // Get tools from message or use registered tools
            auto* toolsNode = msg.data->getChildReadOnly("tools");
            if (toolsNode) {
                // Custom tools from message
                // (would need to parse from IDataNode)
            }

            if (!req.query.empty()) {
                std::lock_guard<std::mutex> lock(m_requestMutex);
                m_requestQueue.push(std::move(req));
                m_requestCV.notify_one();
                m_logger->debug("Request queued: {}", req.query.substr(0, 50));
            }
        }
    }
}

void LLMService::publishResponses() {
    if (!m_io) return;

    std::lock_guard<std::mutex> lock(m_responseMutex);
    while (!m_responseQueue.empty()) {
        auto resp = std::move(m_responseQueue.front());
        m_responseQueue.pop();

        if (resp.isError) {
            auto event = std::make_unique<grove::JsonDataNode>("error");
            event->setString("message", resp.text);
            event->setString("conversationId", resp.conversationId);
            m_io->publish("llm:error", std::move(event));
        } else {
            auto event = std::make_unique<grove::JsonDataNode>("response");
            event->setString("text", resp.text);
            event->setString("conversationId", resp.conversationId);
            event->setInt("tokens", resp.tokens);
            event->setInt("iterations", resp.iterations);
            m_io->publish("llm:response", std::move(event));

            m_logger->info("Response published: {} chars", resp.text.size());
        }
    }
}

void LLMService::workerLoop() {
    m_logger->debug("Worker thread started");

    while (m_running) {
        Request req;

        {
            std::unique_lock<std::mutex> lock(m_requestMutex);
            m_requestCV.wait_for(lock, std::chrono::milliseconds(100), [this] {
                return !m_requestQueue.empty() || !m_running;
            });

            if (!m_running) break;
            if (m_requestQueue.empty()) continue;

            req = std::move(m_requestQueue.front());
            m_requestQueue.pop();
        }

        // Process request (HTTP calls happen here)
        auto resp = processRequest(req);

        {
            std::lock_guard<std::mutex> lock(m_responseMutex);
            m_responseQueue.push(std::move(resp));
        }
    }

    m_logger->debug("Worker thread stopped");
}

LLMService::Response LLMService::processRequest(const Request& request) {
    Response resp;
    resp.conversationId = request.conversationId;

    if (!m_provider) {
        resp.text = "LLM provider not initialized";
        resp.isError = true;
        return resp;
    }

    try {
        // Get or create conversation history
        auto& history = m_conversations[request.conversationId];
        if (history.is_null()) {
            history = nlohmann::json::array();
        }

        // Add user message
        history.push_back({{"role", "user"}, {"content", request.query}});

        // Get tool definitions
        nlohmann::json tools = m_toolRegistry.getToolDefinitions();

        // Run agentic loop
        auto result = agenticLoop(request.query, request.systemPrompt,
                                  history, tools, request.maxIterations);

        if (result.contains("error")) {
            resp.text = result["error"].get<std::string>();
            resp.isError = true;
        } else {
            resp.text = result["response"].get<std::string>();
            resp.tokens = result.value("tokens", 0);
            resp.iterations = result.value("iterations", 1);

            // Add assistant response to history
            history.push_back({{"role", "assistant"}, {"content", resp.text}});
        }

    } catch (const std::exception& e) {
        resp.text = e.what();
        resp.isError = true;
    }

    return resp;
}

nlohmann::json LLMService::agenticLoop(const std::string& query, const std::string& systemPrompt,
                                        nlohmann::json& messages, const nlohmann::json& tools,
                                        int maxIterations) {
    int totalTokens = 0;

    for (int iteration = 0; iteration < maxIterations; iteration++) {
        m_logger->debug("Agentic loop iteration {}", iteration + 1);

        auto response = m_provider->chat(systemPrompt, messages, tools);
        totalTokens += response.input_tokens + response.output_tokens;

        if (response.is_end_turn) {
            return {
                {"response", response.text},
                {"iterations", iteration + 1},
                {"tokens", totalTokens}
            };
        }

        // Execute tool calls
        if (!response.tool_calls.empty()) {
            std::vector<ToolResult> results;

            for (const auto& call : response.tool_calls) {
                m_logger->debug("Executing tool: {}", call.name);
                nlohmann::json result = m_toolRegistry.execute(call.name, call.input);
                results.push_back({call.id, result.dump(), false});
            }

            // Append assistant message and tool results
            m_provider->appendAssistantMessage(messages, response);
            auto toolResultsMsg = m_provider->formatToolResults(results);

            if (toolResultsMsg.is_array()) {
                for (const auto& msg : toolResultsMsg) {
                    messages.push_back(msg);
                }
            } else {
                messages.push_back(toolResultsMsg);
            }
        }
    }

    return {{"error", "max_iterations_reached"}};
}

void LLMService::shutdown() {
    m_running = false;
    m_requestCV.notify_all();

    if (m_workerThread.joinable()) {
        m_workerThread.join();
    }

    m_logger->info("LLMService shutdown");
}

LLMService::SyncResponse LLMService::sendMessageSync(
    const std::string& message,
    const std::string& conversationId,
    const std::string& systemPrompt
) {
    SyncResponse syncResp;

    // Create request (same as async mode)
    Request request;
    request.query = message;
    request.conversationId = conversationId.empty() ? "mcp-session" : conversationId;
    request.systemPrompt = systemPrompt.empty() ? m_defaultSystemPrompt : systemPrompt;
    request.maxIterations = m_maxIterations;

    // Process synchronously (blocking call)
    auto response = processRequest(request);

    // Convert to SyncResponse
    if (!response.isError) {
        syncResp.text = response.text;
        syncResp.tokens = response.tokens;
        syncResp.iterations = response.iterations;
    } else {
        // On error, return error in text
        syncResp.text = "Error: " + response.text;
        syncResp.tokens = 0;
        syncResp.iterations = 0;
    }

    return syncResp;
}

} // namespace celuna
