#include "Mcp_Server.hpp"
#include <mcp/core/Control_Service.hpp>
#include <drogon/drogon.h>
#include <nlohmann/json.hpp>
#include <algorithm>
#include <memory>
#include <string>
#include <thread>

namespace aethera::mcp {
namespace {

constexpr std::string_view protocol_version{"2026-07-28"};

[[nodiscard]] drogon::HttpResponsePtr json_response(
    nlohmann::json value,
    drogon::HttpStatusCode status = drogon::k200OK) {
    auto response = drogon::HttpResponse::newHttpResponse();
    response->setContentTypeCode(drogon::CT_APPLICATION_JSON);
    response->setStatusCode(status);
    response->addHeader("MCP-Protocol-Version", std::string{protocol_version});
    response->setBody(value.dump());
    return response;
}

[[nodiscard]] nlohmann::json rpc_error(
    const nlohmann::json& id, int code, std::string message) {
    return {{"jsonrpc", "2.0"}, {"id", id},
            {"error", {{"code", code}, {"message", std::move(message)}}}};
}

[[nodiscard]] bool accepted_origin(const drogon::HttpRequestPtr& request) {
    const auto origin = request->getHeader("Origin");
    if (origin.empty()) return true;
    return origin == "null" || origin.starts_with("http://127.0.0.1") ||
           origin.starts_with("http://localhost") ||
           origin.starts_with("https://127.0.0.1") ||
           origin.starts_with("https://localhost");
}

[[nodiscard]] nlohmann::json tool_result(const Tool_Call_Output& output) {
    const bool failed = output.result != Tool_Call_Result::ok;
    const auto text = failed ? output.message : output.content.dump();
    nlohmann::json result{
        {"content", nlohmann::json::array({{{"type", "text"}, {"text", text}}})},
        {"isError", failed}};
    if (!output.content.is_null() && !output.content.empty())
        result["structuredContent"] = output.content;
    return result;
}

[[nodiscard]] nlohmann::json dispatch_request(
    Control_Service& control, const nlohmann::json& request) {
    const auto id = request.value("id", nlohmann::json{});
    const auto method = request.value("method", std::string{});
    if (method == "server/discover" || method == "initialize") {
        return {{"jsonrpc", "2.0"}, {"id", id},
                {"result", {
                    {"protocolVersion", protocol_version},
                    {"capabilities", {{"tools", nlohmann::json::object()}}},
                    {"serverInfo", {{"name", "aethera"}, {"version", "1.0.0"}}}}}};
    }
    if (method == "ping")
        return {{"jsonrpc", "2.0"}, {"id", id},
                {"result", nlohmann::json::object()}};
    if (method == "tools/list")
        return {{"jsonrpc", "2.0"}, {"id", id},
                {"result", {{"tools", control.tool_catalog()},
                            {"ttlMs", 1000}, {"cacheScope", "server"}}}};
    if (method == "tools/call") {
        if (!request.contains("params") || !request["params"].is_object())
            return rpc_error(id, -32602, "tools/call requires params");
        const auto& parameters = request["params"];
        const auto name = parameters.value("name", std::string{});
        if (name.empty()) return rpc_error(id, -32602, "tool name is required");
        const auto arguments = parameters.value(
            "arguments", nlohmann::json::object());
        const auto output = control.call_tool(name, arguments);
        if (output.result == Tool_Call_Result::unknown_tool)
            return rpc_error(id, -32602, output.message);
        return {{"jsonrpc", "2.0"}, {"id", id},
                {"result", tool_result(output)}};
    }
    return rpc_error(id, -32601, "method not found");
}

}

int run_mcp_server(std::uint16_t port) {
    auto control = Control_Service::create();
    auto& app = drogon::app();
    app.registerHandler(
        "/mcp",
        [control](const drogon::HttpRequestPtr& request,
                  std::function<void(const drogon::HttpResponsePtr&)>&& callback) {
            if (!accepted_origin(request)) {
                callback(json_response(
                    rpc_error(nullptr, -32000, "origin is not allowed"),
                    drogon::k403Forbidden));
                return;
            }
            if (request->method() == drogon::Get) {
                callback(json_response(
                    rpc_error(nullptr, -32600, "SSE stream is not provided"),
                    drogon::k405MethodNotAllowed));
                return;
            }
            nlohmann::json message;
            try {
                message = nlohmann::json::parse(request->body());
            } catch (const nlohmann::json::exception&) {
                callback(json_response(
                    rpc_error(nullptr, -32700, "invalid JSON"),
                    drogon::k400BadRequest));
                return;
            }
            if (!message.is_object() || message.value("jsonrpc", "") != "2.0" ||
                !message.contains("method")) {
                callback(json_response(
                    rpc_error(message.value("id", nlohmann::json{}),
                              -32600, "invalid JSON-RPC request"),
                    drogon::k400BadRequest));
                return;
            }
            if (!message.contains("id")) {
                auto response = drogon::HttpResponse::newHttpResponse();
                response->setStatusCode(drogon::k202Accepted);
                callback(std::move(response));
                return;
            }
            try {
                callback(json_response(dispatch_request(*control, message)));
            } catch (const std::exception& failure) {
                callback(json_response(
                    rpc_error(message["id"], -32603, failure.what()),
                    drogon::k500InternalServerError));
            }
        },
        {drogon::Get, drogon::Post});
    app.addListener("127.0.0.1", port)
       .setThreadNum(std::min(8U, std::max(2U, std::thread::hardware_concurrency())))
       .setIdleConnectionTimeout(90)
       .run();
    return 0;
}

}