/**
 * @file host_example.cpp
 * @brief Example MCP host implementation combining server and client capabilities
 */

#include "mcp_server.h"
#include "mcp_client.h"
#include "mcp_resource.h"
#include "mcp_tools.h"
#include "mcp_workflow_resource.h"
#include "custom_agent.h"

#include <iostream>
#include <string>
#include <vector>
#include <map>
#include <chrono>
#include <thread>
#include <functional>
#include <mutex>
#include <condition_variable>
#include <atomic>

/**
 * @class mcp_host
 * @brief A host that combines MCP server and client capabilities
 * 
 * This class provides an example of a host application that can act both
 * as an MCP server (exposing resources and tools) and an MCP client
 * (connecting to other MCP servers).
 */
class mcp_host {
public:
    /**
     * @brief Constructor
     * @param name Host name
     * @param host Host address to bind to
     * @param port Port to listen on
     */
    mcp_host(const std::string& name, const std::string& host = "localhost", int port = 8080)
        : name_(name), running_(false) {
        
        // Create the server
        server_ = std::make_unique<mcp::server>(host, port);
        server_->set_name(name);
        
        // Set up workflow resource
        workflow_resource_ = std::make_shared<mcp::workflow_resource>();
        server_->register_resource("/workflows", workflow_resource_);
        
        // Set up agent resource
        agent_resource_ = std::make_shared<mcp::agent_resource>();
        server_->register_resource("/agents", agent_resource_);
        
        // Register host management API
        host_api_ = std::make_shared<mcp::api_resource>("Host API", "Host management API");
        
        // Register GET /status endpoint
        host_api_->register_handler(
            mcp::http_method::get,
            "/status",
            [this](const mcp::request& req) -> mcp::response {
                mcp::response res;
                res.set_json_body({
                    {"name", name_},
                    {"running", running_.load()},
                    {"connections", connections_.size()},
                    {"agents", agent_resource_->get_agent_names()},
                    {"workflows", workflow_resource_->get_workflow_names()}
                });
                return res;
            },
            "Get host status"
        );
        
        // Register POST /connect endpoint
        host_api_->register_handler(
            mcp::http_method::post,
            "/connect",
            [this](const mcp::request& req) -> mcp::response {
                mcp::response res;
                try {
                    mcp::json body = req.json_body();
                    
                    if (!body.contains("host") || !body["host"].is_string() ||
                        !body.contains("port") || !body["port"].is_number_integer()) {
                        res.set_error(mcp::error_code::invalid_request, 
                                     "Missing required fields: host (string) and port (integer)");
                        return res;
                    }
                    
                    std::string host = body["host"];
                    int port = body["port"];
                    
                    std::string connection_id = connect_to_server(host, port);
                    
                    res.set_json_body({
                        {"connection_id", connection_id},
                        {"status", "connected"}
                    });
                    
                } catch (const std::exception& e) {
                    res.set_error(mcp::error_code::internal_server_error, e.what());
                }
                return res;
            },
            "Connect to another MCP server"
        );
        
        // Register GET /connections endpoint
        host_api_->register_handler(
            mcp::http_method::get,
            "/connections",
            [this](const mcp::request& req) -> mcp::response {
                mcp::response res;
                mcp::json connections = mcp::json::array();
                
                std::lock_guard<std::mutex> lock(connections_mutex_);
                for (const auto& [id, conn] : connections_) {
                    connections.push_back({
                        {"id", id},
                        {"host", conn.host},
                        {"port", conn.port}
                    });
                }
                
                res.set_json_body(connections);
                return res;
            },
            "List all connections"
        );
        
        // Register DELETE /connections/{id} endpoint
        host_api_->register_handler(
            mcp::http_method::delete_,
            "/connections/:id",
            [this](const mcp::request& req) -> mcp::response {
                mcp::response res;
                
                // Extract connection ID from path
                std::string path = req.path;
                size_t pos = path.find_last_of('/');
                if (pos == std::string::npos) {
                    res.set_error(mcp::error_code::invalid_request, "Invalid path");
                    return res;
                }
                
                std::string connection_id = path.substr(pos + 1);
                
                bool success = disconnect_from_server(connection_id);
                
                if (success) {
                    res.status_code = 204;  // No content
                } else {
                    res.set_error(mcp::error_code::not_found, 
                                 "Connection not found: " + connection_id);
                }
                
                return res;
            },
            "Disconnect from a server"
        );
        
        server_->register_resource("/host", host_api_);
    }
    
    /**
     * @brief Destructor
     */
    ~mcp_host() {
        stop();
    }
    
    /**
     * @brief Start the host
     * @param blocking If true, this call blocks until the host stops
     */
    void start(bool blocking = false) {
        if (running_) {
            return;
        }
        
        running_ = true;
        
        // Start the server in a separate thread
        server_thread_ = std::thread([this]() {
            try {
                server_->start(true);  // Blocking call
            } catch (const std::exception& e) {
                std::cerr << "Server error: " << e.what() << std::endl;
                running_ = false;
            }
        });
        
        if (blocking) {
            // Wait for the server to exit
            if (server_thread_.joinable()) {
                server_thread_.join();
            }
        }
    }
    
    /**
     * @brief Stop the host
     */
    void stop() {
        if (!running_) {
            return;
        }
        
        running_ = false;
        
        // Disconnect from all servers
        {
            std::lock_guard<std::mutex> lock(connections_mutex_);
            connections_.clear();
        }
        
        // Stop the server
        server_->stop();
        
        // Wait for the server thread to exit
        if (server_thread_.joinable()) {
            server_thread_.join();
        }
    }
    
    /**
     * @brief Register a tool
     * @param tool The tool to register
     * @param handler The function to call when the tool is invoked
     */
    void register_tool(const mcp::tool& tool, mcp::tool_handler handler) {
        server_->register_tool(tool, handler);
    }
    
    /**
     * @brief Register a resource
     * @param path The path to mount the resource at
     * @param resource The resource to register
     */
    void register_resource(const std::string& path, std::shared_ptr<mcp::resource> resource) {
        server_->register_resource(path, resource);
    }
    
    /**
     * @brief Register an agent
     * @param agent The agent to register
     */
    void register_agent(std::shared_ptr<mcp::agent> agent) {
        agent_resource_->register_agent(agent);
    }
    
    /**
     * @brief Register a workflow
     * @param workflow The workflow to register
     */
    void register_workflow(const mcp::workflow& workflow) {
        workflow_resource_->register_workflow(workflow);
    }
    
    /**
     * @brief Connect to another MCP server
     * @param host The server host
     * @param port The server port
     * @return Connection ID
     */
    std::string connect_to_server(const std::string& host, int port) {
        // Create a unique connection ID
        static int next_id = 1;
        std::string connection_id = "conn_" + std::to_string(next_id++);
        
        // Create the client
        auto client = std::make_shared<mcp::client>(host, port);
        
        // Test the connection by getting server info
        try {
            mcp::json server_info = client->get_server_info();
            std::cout << "Connected to server: " << server_info["name"].get<std::string>() 
                      << " (version " << server_info["version"].get<std::string>() << ")"
                      << std::endl;
        } catch (const std::exception& e) {
            throw std::runtime_error(
                "Failed to connect to server at " + host + ":" + std::to_string(port) + 
                " - " + e.what()
            );
        }
        
        // Store the connection
        {
            std::lock_guard<std::mutex> lock(connections_mutex_);
            connections_[connection_id] = {host, port, client};
        }
        
        return connection_id;
    }
    
    /**
     * @brief Disconnect from a server
     * @param connection_id The connection ID
     * @return True if the connection was found and removed
     */
    bool disconnect_from_server(const std::string& connection_id) {
        std::lock_guard<std::mutex> lock(connections_mutex_);
        return connections_.erase(connection_id) > 0;
    }
    
    /**
     * @brief Get a client by connection ID
     * @param connection_id The connection ID
     * @return The client, or nullptr if not found
     */
    std::shared_ptr<mcp::client> get_client(const std::string& connection_id) {
        std::lock_guard<std::mutex> lock(connections_mutex_);
        auto it = connections_.find(connection_id);
        if (it != connections_.end()) {
            return it->second.client;
        }
        return nullptr;
    }
    
    /**
     * @brief Call a tool on a connected server
     * @param connection_id The connection ID
     * @param tool_name The name of the tool to call
     * @param parameters The parameters to pass to the tool
     * @return The result of the tool call
     */
    mcp::json call_remote_tool(const std::string& connection_id, 
                              const std::string& tool_name, 
                              const mcp::json& parameters = mcp::json::object()) {
        auto client = get_client(connection_id);
        if (!client) {
            throw std::runtime_error("Connection not found: " + connection_id);
        }
        
        return client->call_tool(tool_name, parameters);
    }

private:
    struct connection {
        std::string host;
        int port;
        std::shared_ptr<mcp::client> client;
    };
    
    std::string name_;
    std::atomic<bool> running_;
    
    std::unique_ptr<mcp::server> server_;
    std::thread server_thread_;
    
    std::shared_ptr<mcp::workflow_resource> workflow_resource_;
    std::shared_ptr<mcp::agent_resource> agent_resource_;
    std::shared_ptr<mcp::api_resource> host_api_;
    
    std::map<std::string, connection> connections_;
    std::mutex connections_mutex_;
};

// Example tool handler for system information
mcp::json system_info_handler(const mcp::json& params) {
    mcp::json info = {
        {"hostname", "example-host"},
        {"platform", "example-platform"},
        {"uptime_seconds", 3600}
    };
    
    // This would be populated with actual system info in a real implementation
    
    return info;
}

int main() {
    // Create the host
    mcp_host host("Example MCP Host", "localhost", 8081);
    
    // Register a system info tool
    mcp::tool system_info_tool = mcp::tool_builder("system_info")
        .with_description("Get system information")
        .build();
    
    host.register_tool(system_info_tool, system_info_handler);
    
    // Create an example file resource
    auto file_resource = std::make_shared<mcp::file_resource>("./files");
    host.register_resource("/files", file_resource);
    
    // Create a chain agent for processing
    auto chain_agent_ptr = std::make_shared<examples::chain_agent>("processor");
    
    // Set up a chain that calls system_info, then does text processing
    chain_agent_ptr->initialize({
        {"chain", mcp::json::array({
            {
                {"tool", "system_info"}
            },
            {
                {"tool", "text_processor"},
                {"mappings", {
                    {"hostname", "text"}  // Map system_info.hostname to text_processor.text
                }}
            }
        })}
    });
    
    host.register_agent(chain_agent_ptr);
    
    // Create a workflow that combines local and remote tools (when connected)
    mcp::workflow combined_workflow("combined", "Combined workflow using local and remote tools");
    combined_workflow.add_tool_call("system_info");
    
    host.register_workflow(combined_workflow);
    
    std::cout << "Starting Host on http://localhost:8081" << std::endl;
    std::cout << "The host combines server and client capabilities." << std::endl;
    std::cout << "You can interact with it via its API at /host" << std::endl;
    std::cout << "Press Ctrl+C to stop" << std::endl;
    
    try {
        // Start the host in blocking mode
        host.start(true);
    } catch (const std::exception& e) {
        std::cerr << "Error: " << e.what() << std::endl;
        return 1;
    }
    
    return 0;
}