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

#include <iostream>
#include <sstream>

namespace celuna::mcp {

MCPServer::MCPServer(celuna::ToolRegistry& registry) : m_registry(registry) {
    m_logger = spdlog::get("MCPServer");
    if (!m_logger) {
        // Log to stderr so stdout stays clean for JSON-RPC
        m_logger = spdlog::stderr_color_mt("MCPServer");
    }

    // Default capabilities
    m_capabilities = {
        {"tools", {{"listChanged", false}}}
    };
}

MCPServer::~MCPServer() {
    stop();
}

void MCPServer::setServerInfo(const std::string& name, const std::string& version) {
    m_serverName = name;
    m_serverVersion = version;
}

void MCPServer::setCapabilities(const json& capabilities) {
    m_capabilities = capabilities;
}

void MCPServer::run() {
    m_running = true;
    m_logger->info("MCP Server starting (stdio mode)");

    while (m_running) {
        std::string line = readLine();
        if (line.empty()) {
            if (m_running) {
                m_logger->debug("EOF received, stopping server");
            }
            break;
        }

        processLine(line);
    }

    m_logger->info("MCP Server stopped");
}

void MCPServer::stop() {
    m_running = false;
}

std::string MCPServer::readLine() {
    std::string line;
    if (!std::getline(std::cin, line)) {
        return "";
    }
    // Remove trailing \r if present (Windows line endings)
    if (!line.empty() && line.back() == '\r') {
        line.pop_back();
    }
    return line;
}

void MCPServer::writeLine(const std::string& line) {
    std::cout << line << std::endl;
    std::cout.flush();
}

void MCPServer::processLine(const std::string& line) {
    m_logger->debug("Received: {}", line.substr(0, 200));

    try {
        json message = json::parse(line);

        // Validate JSON-RPC
        if (!message.contains("jsonrpc") || message["jsonrpc"] != "2.0") {
            sendError(-1, -32600, "Invalid JSON-RPC: missing or invalid jsonrpc field");
            return;
        }

        // Check if it's a request (has method) or notification (no id)
        if (!message.contains("method")) {
            sendError(message.value("id", -1), -32600, "Invalid Request: missing method");
            return;
        }

        JsonRpcRequest request;
        request.method = message["method"].get<std::string>();
        request.params = message.value("params", json::object());
        request.id = message.value("id", 0);

        bool isNotification = !message.contains("id");

        json result = handleRequest(request);

        // Only send response for requests, not notifications
        if (!isNotification) {
            if (result.contains("error")) {
                sendError(request.id, result["error"]["code"].get<int>(),
                          result["error"]["message"].get<std::string>());
            } else {
                sendResponse(request.id, result);
            }
        }

    } catch (const json::exception& e) {
        m_logger->error("JSON parse error: {}", e.what());
        sendError(-1, -32700, std::string("Parse error: ") + e.what());
    }
}

json MCPServer::handleRequest(const JsonRpcRequest& request) {
    m_logger->debug("Handling method: {}", request.method);

    // MCP methods
    if (request.method == "initialize") {
        return handleInitialize(request.params);
    } else if (request.method == "ping") {
        return handlePing(request.params);
    } else if (request.method == "tools/list") {
        return handleToolsList(request.params);
    } else if (request.method == "tools/call") {
        return handleToolsCall(request.params);
    } else if (request.method == "notifications/initialized") {
        // Client notification that initialization is complete
        m_logger->info("Client initialized");
        return json::object();  // No response for notifications
    }

    // Unknown method
    return {{"error", {{"code", -32601}, {"message", "Method not found: " + request.method}}}};
}

json MCPServer::handleInitialize(const json& params) {
    if (m_initialized) {
        return {{"error", {{"code", -32600}, {"message", "Already initialized"}}}};
    }

    // Extract client info
    if (params.contains("clientInfo")) {
        m_clientName = params["clientInfo"].value("name", "unknown");
        m_clientVersion = params["clientInfo"].value("version", "unknown");
        m_logger->info("Client connected: {} v{}", m_clientName, m_clientVersion);
    }

    m_initialized = true;

    // Return server info and capabilities
    return {
        {"protocolVersion", "2024-11-05"},
        {"serverInfo", {
            {"name", m_serverName},
            {"version", m_serverVersion}
        }},
        {"capabilities", m_capabilities}
    };
}

json MCPServer::handleToolsList(const json& params) {
    if (!m_initialized) {
        return {{"error", {{"code", -32600}, {"message", "Not initialized"}}}};
    }

    // Get tool definitions from registry and convert to MCP format
    json registryTools = m_registry.getToolDefinitions();
    json tools = json::array();

    for (const auto& tool : registryTools) {
        json toolEntry;
        toolEntry["name"] = tool["name"];
        toolEntry["description"] = tool["description"];
        toolEntry["inputSchema"] = tool["input_schema"];
        tools.push_back(toolEntry);
    }

    return {{"tools", tools}};
}

json MCPServer::handleToolsCall(const json& params) {
    if (!m_initialized) {
        return {{"error", {{"code", -32600}, {"message", "Not initialized"}}}};
    }

    if (!params.contains("name")) {
        return {{"error", {{"code", -32602}, {"message", "Missing tool name"}}}};
    }

    std::string toolName = params["name"].get<std::string>();
    json arguments = params.value("arguments", json::object());

    m_logger->info("Calling tool: {}", toolName);

    try {
        json result = m_registry.execute(toolName, arguments);

        // Format result as MCP content
        json content = json::array();

        // Convert result to text content
        std::string textContent;
        if (result.is_string()) {
            textContent = result.get<std::string>();
        } else {
            textContent = result.dump(2);
        }

        json contentEntry;
        contentEntry["type"] = "text";
        contentEntry["text"] = textContent;
        content.push_back(contentEntry);

        return {
            {"content", content},
            {"isError", false}
        };

    } catch (const std::exception& e) {
        m_logger->error("Tool execution failed: {}", e.what());
        json errorContent = json::array();
        json errorEntry;
        errorEntry["type"] = "text";
        errorEntry["text"] = e.what();
        errorContent.push_back(errorEntry);

        json errorResult;
        errorResult["content"] = errorContent;
        errorResult["isError"] = true;
        return errorResult;
    }
}

json MCPServer::handlePing(const json& params) {
    return json::object();  // Empty response = pong
}

void MCPServer::sendResponse(int id, const json& result) {
    json response = {
        {"jsonrpc", "2.0"},
        {"id", id},
        {"result", result}
    };
    std::string line = response.dump();
    m_logger->debug("Sending: {}", line.substr(0, 200));
    writeLine(line);
}

void MCPServer::sendError(int id, int code, const std::string& message) {
    json response = {
        {"jsonrpc", "2.0"},
        {"id", id},
        {"error", {
            {"code", code},
            {"message", message}
        }}
    };
    writeLine(response.dump());
}

void MCPServer::sendNotification(const std::string& method, const json& params) {
    json notification = {
        {"jsonrpc", "2.0"},
        {"method", method}
    };
    if (!params.empty()) {
        notification["params"] = params;
    }
    writeLine(notification.dump());
}

} // namespace celuna::mcp
