#include "MCPClient.hpp"
#include <spdlog/sinks/stdout_color_sinks.h>

#include <fstream>

namespace celuna::mcp {

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

MCPClient::~MCPClient() {
    disconnectAll();
}

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

        json config;
        file >> config;

        if (!config.contains("servers")) {
            m_logger->warn("No 'servers' key in MCP config");
            return false;
        }

        m_configs.clear();
        for (auto& [name, serverConfig] : config["servers"].items()) {
            auto cfg = MCPServerConfig::fromJson(name, serverConfig);
            if (cfg.enabled && !cfg.command.empty()) {
                m_configs[name] = cfg;
                m_logger->info("Loaded MCP server config: {}", name);
            }
        }

        m_logger->info("Loaded {} MCP server configs", m_configs.size());
        return true;

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

int MCPClient::connectAll() {
    int connected = 0;

    for (const auto& [name, config] : m_configs) {
        if (connect(name)) {
            connected++;
        }
    }

    m_logger->info("Connected to {}/{} MCP servers", connected, m_configs.size());
    return connected;
}

bool MCPClient::connect(const std::string& serverName) {
    auto configIt = m_configs.find(serverName);
    if (configIt == m_configs.end()) {
        m_logger->error("Unknown MCP server: {}", serverName);
        return false;
    }

    if (m_servers.count(serverName) > 0 && m_servers[serverName].initialized) {
        m_logger->debug("Server {} already connected", serverName);
        return true;
    }

    m_logger->info("Connecting to MCP server: {}", serverName);

    ConnectedServer server;
    server.config = configIt->second;
    server.transport = std::make_unique<StdioTransport>(server.config);

    if (!server.transport->start()) {
        m_logger->error("Failed to start MCP server: {}", serverName);
        return false;
    }

    if (!initializeServer(server)) {
        m_logger->error("Failed to initialize MCP server: {}", serverName);
        server.transport->stop();
        return false;
    }

    if (!discoverTools(server)) {
        m_logger->warn("Failed to discover tools for: {}", serverName);
        // Continue anyway, server might not have tools
    }

    m_servers[serverName] = std::move(server);
    m_logger->info("Connected to MCP server: {} ({} tools)",
                   serverName, m_servers[serverName].tools.size());

    return true;
}

void MCPClient::disconnect(const std::string& serverName) {
    auto it = m_servers.find(serverName);
    if (it == m_servers.end()) {
        return;
    }

    if (it->second.transport) {
        it->second.transport->stop();
    }

    m_servers.erase(it);
    m_logger->info("Disconnected from MCP server: {}", serverName);
}

void MCPClient::disconnectAll() {
    for (auto& [name, server] : m_servers) {
        if (server.transport) {
            server.transport->stop();
        }
    }
    m_servers.clear();
}

bool MCPClient::initializeServer(ConnectedServer& server) {
    // Send initialize request
    JsonRpcRequest initReq;
    initReq.method = "initialize";
    initReq.params = {
        {"protocolVersion", "2024-11-05"},
        {"capabilities", {
            {"roots", {{"listChanged", true}}}
        }},
        {"clientInfo", {
            {"name", "AISSIA"},
            {"version", "1.0.0"}
        }}
    };

    auto response = server.transport->sendRequest(initReq, 10000);

    if (response.isError()) {
        m_logger->error("Initialize failed: {}",
                        response.error->dump());
        return false;
    }

    if (response.result.has_value()) {
        server.info = MCPServerInfo::fromJson(*response.result);
        m_logger->debug("Server info: {} v{}",
                        server.info.name, server.info.version);
    }

    // Send initialized notification
    server.transport->sendNotification("notifications/initialized", json::object());

    server.initialized = true;
    return true;
}

bool MCPClient::discoverTools(ConnectedServer& server) {
    if (!server.info.capabilities.hasTools) {
        m_logger->debug("Server {} does not support tools", server.config.name);
        return true;
    }

    JsonRpcRequest listReq;
    listReq.method = "tools/list";
    listReq.params = json::object();

    auto response = server.transport->sendRequest(listReq, 10000);

    if (response.isError()) {
        m_logger->error("tools/list failed: {}",
                        response.error->dump());
        return false;
    }

    server.tools.clear();

    if (response.result.has_value() && response.result->contains("tools")) {
        for (const auto& toolJson : (*response.result)["tools"]) {
            MCPTool tool = MCPTool::fromJson(toolJson);
            server.tools.push_back(tool);
            m_logger->debug("Discovered tool: {}:{}", server.config.name, tool.name);
        }
    }

    return true;
}

std::vector<MCPTool> MCPClient::listAllTools() {
    std::vector<MCPTool> allTools;

    for (const auto& [serverName, server] : m_servers) {
        for (const auto& tool : server.tools) {
            MCPTool prefixedTool = tool;
            prefixedTool.name = serverName + ":" + tool.name;
            allTools.push_back(prefixedTool);
        }
    }

    return allTools;
}

MCPToolResult MCPClient::callTool(const std::string& toolName, const json& arguments) {
    auto [serverName, actualToolName] = parseToolName(toolName);

    auto it = m_servers.find(serverName);
    if (it == m_servers.end()) {
        MCPToolResult error;
        error.isError = true;
        error.content = {{{"type", "text"}, {"text", "Server not connected: " + serverName}}};
        return error;
    }

    JsonRpcRequest callReq;
    callReq.method = "tools/call";
    callReq.params = {
        {"name", actualToolName},
        {"arguments", arguments}
    };

    m_logger->debug("Calling tool {}:{} with args: {}",
                    serverName, actualToolName, arguments.dump().substr(0, 100));

    auto response = it->second.transport->sendRequest(callReq, 30000);

    MCPToolResult result;

    if (response.isError()) {
        result.isError = true;
        std::string errorMsg = response.error.has_value() ?
            response.error->dump() : "Unknown error";
        result.content = {{{"type", "text"}, {"text", errorMsg}}};
        m_logger->error("Tool call failed: {}", errorMsg);
        return result;
    }

    if (response.result.has_value()) {
        auto& res = *response.result;
        if (res.contains("content")) {
            for (const auto& contentItem : res["content"]) {
                result.content.push_back(contentItem);
            }
        }
        result.isError = res.value("isError", false);
    }

    m_logger->debug("Tool {} returned {} content items",
                    toolName, result.content.size());

    return result;
}

size_t MCPClient::toolCount() const {
    size_t count = 0;
    for (const auto& [_, server] : m_servers) {
        count += server.tools.size();
    }
    return count;
}

bool MCPClient::isConnected(const std::string& serverName) const {
    auto it = m_servers.find(serverName);
    return it != m_servers.end() && it->second.initialized;
}

std::vector<std::string> MCPClient::getConnectedServers() const {
    std::vector<std::string> names;
    for (const auto& [name, _] : m_servers) {
        names.push_back(name);
    }
    return names;
}

std::pair<std::string, std::string> MCPClient::parseToolName(const std::string& toolName) {
    size_t colonPos = toolName.find(':');
    if (colonPos == std::string::npos) {
        // No prefix, try to find tool in any server
        for (const auto& [serverName, server] : m_servers) {
            for (const auto& tool : server.tools) {
                if (tool.name == toolName) {
                    return {serverName, toolName};
                }
            }
        }
        return {"", toolName};
    }

    return {
        toolName.substr(0, colonPos),
        toolName.substr(colonPos + 1)
    };
}

} // namespace celuna::mcp
