/**
 * @file mcp_client.cpp
 * @brief Implementation of the MCP client
 * 
 * This file implements the client-side functionality for the Model Context Protocol.
 * Follows the 2024-11-05 basic protocol specification.
 */

#include "mcp_client.h"
#include "base64.hpp"

namespace mcp {

client::client(const std::string& host, int port, const json& capabilities)
    : host_(host), port_(port), capabilities_(capabilities) {
    
    init_client();
}

client::~client() {
    // httplib::Client will be automatically destroyed
}

void client::init_client() {
    // Create the HTTP client
    http_client_ = std::make_unique<httplib::Client>(host_.c_str(), port_);

    // Set timeout
    http_client_->set_connection_timeout(timeout_seconds_, 0);
    http_client_->set_read_timeout(timeout_seconds_, 0);
    http_client_->set_write_timeout(timeout_seconds_, 0);
}

bool client::initialize(const std::string& client_name, const std::string& client_version) {
    // Create initialization request
    request req = request::create("initialize", {
        {"protocolVersion", MCP_VERSION},
        {"capabilities", capabilities_},
        {"clientInfo", {
            {"name", client_name},
            {"version", client_version}
        }}
    });
    
    try {
        // Send the request
        json result = send_jsonrpc(req);
        
        // Store server capabilities
        server_capabilities_ = result["capabilities"];
        
        // Send initialized notification
        request notification = request::create_notification("initialized");
        send_jsonrpc(notification);
        
        return true;
    } catch (const std::exception& e) {
        // Initialization failed
        return false;
    }
}

bool client::ping() {
    // Create ping request
    request req = request::create("ping", {});
    
    try {
        // Send the request
        json result = send_jsonrpc(req);

        // The receiver MUST respond promptly with an empty response
        if (result.empty()) {
            return true;
        } else {
            return false;
        }
    } catch (const std::exception& e) {
        // Ping failed
        return false;
    }
}

void client::set_auth_token(const std::string& token) {
    std::lock_guard<std::mutex> lock(mutex_);
    auth_token_ = token;
    
    // Add to default headers
    set_header("Authorization", "Bearer " + auth_token_);
}

void client::set_header(const std::string& key, const std::string& value) {
    std::lock_guard<std::mutex> lock(mutex_);
    default_headers_[key] = value;
}

void client::set_timeout(int timeout_seconds) {
    std::lock_guard<std::mutex> lock(mutex_);
    timeout_seconds_ = timeout_seconds;
    
    // Update the client's timeout
    http_client_->set_connection_timeout(timeout_seconds_, 0);
    http_client_->set_read_timeout(timeout_seconds_, 0);
    http_client_->set_write_timeout(timeout_seconds_, 0);
}

void client::set_capabilities(const json& capabilities) {
    std::lock_guard<std::mutex> lock(mutex_);
    capabilities_ = capabilities;
}

response client::send_request(const std::string& method, const json& params) {
    request req = request::create(method, params);
    json result = send_jsonrpc(req);
    
    response res;
    res.jsonrpc = "2.0";
    res.id = req.id;
    res.result = result;
    
    return res;
}

void client::send_notification(const std::string& method, const json& params) {
    request req = request::create_notification(method, params);
    send_jsonrpc(req);
}

json client::get_server_capabilities() {
    return server_capabilities_;
}

json client::call_tool(const std::string& tool_name, const json& arguments) {
    return send_request("tools/call", {
        {"name", tool_name},
        {"arguments", arguments}
    }).result;
}

std::vector<tool> client::get_tools() {
    json tools_json = send_request("tools/list", {}).result;
    std::vector<tool> tools;
    
    if (tools_json.is_array()) {
        for (const auto& tool_json : tools_json) {
            tool t;
            t.name = tool_json["name"];
            t.description = tool_json["description"];
            
            if (tool_json.contains("inputSchema")) {
                t.parameters_schema = tool_json["inputSchema"];
            }
            
            tools.push_back(t);
        }
    }
    
    return tools;
}

json client::get_resource_metadata(const std::string& resource_path) {
    return send_request("resources/metadata", {
        {"path", resource_path}
    }).result;
}

json client::get_capabilities() {
    return capabilities_;
}

json client::access_resource(const std::string& resource_path, const json& params) {
    json request_params = {
        {"path", resource_path}
    };
    
    // Add any additional parameters
    for (auto it = params.begin(); it != params.end(); ++it) {
        request_params[it.key()] = it.value();
    }
    
    return send_request("resources/access", request_params).result;
}

json client::send_jsonrpc(const request& req) {
    std::lock_guard<std::mutex> lock(mutex_);
    
    // Convert request to JSON
    json req_json = req.to_json();
    std::string req_body = req_json.dump();
    
    // Prepare headers
    httplib::Headers headers;
    headers.emplace("Content-Type", "application/json");
    
    // Add default headers
    for (const auto& [key, value] : default_headers_) {
        headers.emplace(key, value);
    }
    
    // Send the request
    auto result = http_client_->Post("/jsonrpc", headers, req_body, "application/json");
    
    if (!result) {
        // Error occurred
        auto err = result.error();
        
        switch (err) {
            case httplib::Error::Connection:
                throw mcp_exception(error_code::server_error_start, "Connection error");
            case httplib::Error::Read:
                throw mcp_exception(error_code::internal_error, "Read error");
            case httplib::Error::Write:
                throw mcp_exception(error_code::internal_error, "Write error");
            case httplib::Error::ConnectionTimeout:
                throw mcp_exception(error_code::server_error_start, "Timeout error");
            default:
                throw mcp_exception(error_code::internal_error, 
                                  "HTTP client error: " + std::to_string(static_cast<int>(err)));
        }
    }
    
    // Check if it's a notification (no response expected)
    if (req.is_notification()) {
        return json::object();
    }
    
    // Parse response
    try {
        json res_json = json::parse(result->body);
        
        // Check for error
        if (res_json.contains("error")) {
            int code = res_json["error"]["code"];
            std::string message = res_json["error"]["message"];
            
            throw mcp_exception(static_cast<error_code>(code), message);
        }
        
        // Return result
        if (res_json.contains("result")) {
            return res_json["result"];
        } else {
            return json::object();
        }
    } catch (const json::exception& e) {
        throw mcp_exception(error_code::parse_error, 
                          "Failed to parse JSON-RPC response: " + std::string(e.what()));
    }
}

} // namespace mcp