/**
 * @file mcp.server.cpp
 * @brief Example MCP server implementation with custom resources and tools
 */

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

#include <iostream>
#include <string>
#include <chrono>
#include <thread>
#include <ctime>
#include <iomanip>
#include <sstream>

// 添加日志记录功能
class Logger {
public:
    enum class LogLevel {
        DEBUG,
        INFO,
        WARNING,
        ERROR
    };
    
    static std::string getCurrentTimestamp() {
        auto now = std::chrono::system_clock::now();
        auto now_ms = std::chrono::time_point_cast<std::chrono::milliseconds>(now);
        auto time_t_now = std::chrono::system_clock::to_time_t(now);
        auto ms = now_ms.time_since_epoch().count() % 1000;
        
        std::stringstream ss;
        ss << std::put_time(std::localtime(&time_t_now), "%Y-%m-%d %H:%M:%S");
        ss << "." << std::setfill('0') << std::setw(3) << ms;
        return ss.str();
    }
    
    static std::string getLevelString(LogLevel level) {
        switch (level) {
            case LogLevel::DEBUG: return "DEBUG   ";
            case LogLevel::INFO: return "INFO    ";
            case LogLevel::WARNING: return "WARNING ";
            case LogLevel::ERROR: return "ERROR   ";
            default: return "UNKNOWN ";
        }
    }
    
    static void log(LogLevel level, const std::string& module, const std::string& message) {
        std::cout << getCurrentTimestamp() << " | " 
                  << getLevelString(level) << " | " 
                  << module << ": " 
                  << message << std::endl;
    }
    
    static void debug(const std::string& module, const std::string& message) {
        log(LogLevel::DEBUG, module, message);
    }
    
    static void info(const std::string& module, const std::string& message) {
        log(LogLevel::INFO, module, message);
    }
    
    static void warning(const std::string& module, const std::string& message) {
        log(LogLevel::WARNING, module, message);
    }
    
    static void error(const std::string& module, const std::string& message) {
        log(LogLevel::ERROR, module, message);
    }
};

// 创建一个日志记录资源包装器
class LoggingResourceWrapper : public mcp::resource {
public:
    LoggingResourceWrapper(std::shared_ptr<mcp::resource> wrapped_resource, const std::string& resource_path)
        : wrapped_resource_(wrapped_resource), resource_path_(resource_path) {}
    
    mcp::resource_type type() const override {
        return wrapped_resource_->type();
    }
    
    mcp::json metadata() const override {
        return wrapped_resource_->metadata();
    }
    
    mcp::response handle_request(const mcp::request& req) override {
        // 记录请求信息
        std::stringstream ss;
        ss << "Received " << mcp::http_method_to_string(req.method) << " request to " 
           << resource_path_ << req.path;
        
        // 添加查询参数
        if (!req.query_params.empty()) {
            ss << " with params: {";
            bool first = true;
            for (const auto& [key, value] : req.query_params) {
                if (!first) ss << ", ";
                ss << key << ": " << value;
                first = false;
            }
            ss << "}";
        }
        
        // 添加请求体信息（如果不是GET请求）
        if (req.method != mcp::http_method::get && !req.body.empty()) {
            if (req.body.length() > 1000) {
                ss << " with body: " << req.body.substr(0, 1000) << "... (truncated)";
            } else {
                ss << " with body: " << req.body;
            }
        }
        
        Logger::info("mcp.server.handle_request:", ss.str());
        
        // 调用包装的资源处理请求
        mcp::response res = wrapped_resource_->handle_request(req);
        
        // 记录响应信息
        std::stringstream res_ss;
        res_ss << "Response status: " << res.status_code;
        if (!res.body.empty()) {
            if (res.body.length() > 1000) {
                res_ss << ", body: " << res.body.substr(0, 1000) << "... (truncated)";
            } else {
                res_ss << ", body: " << res.body;
            }
        }
        
        Logger::info("mcp.server.handle_request:", res_ss.str());
        
        return res;
    }
    
    mcp::json schema() const override {
        return wrapped_resource_->schema();
    }

private:
    std::shared_ptr<mcp::resource> wrapped_resource_;
    std::string resource_path_;
};

// Example tool handler for getting the current time
mcp::json get_time_handler(const mcp::json& params) {
    Logger::info("mcp.server.toolcall:execute", "Executing get_time tool");
    
    auto now = std::chrono::system_clock::now();
    auto time_t_now = std::chrono::system_clock::to_time_t(now);
    
    std::stringstream formatted_time;
    formatted_time << std::put_time(std::localtime(&time_t_now), "%Y-%m-%d %H:%M:%S");
    
    mcp::json result = {
        {"timestamp", static_cast<int>(time_t_now)},
        {"formatted", formatted_time.str()}
    };
    
    Logger::info("mcp.server.toolcall:result", "get_time result: " + result.dump());
    
    return result;
}

// Example tool handler for mathematical calculations
mcp::json calculator_handler(const mcp::json& params) {
    Logger::info("mcp.server.toolcall:execute", "Executing calculator tool with params: " + params.dump());
    
    if (!params.contains("operation") || !params["operation"].is_string()) {
        Logger::error("mcp.server.toolcall:execute", "Missing required parameter: operation (string)");
        throw std::runtime_error("Missing required parameter: operation (string)");
    }
    
    std::string op = params["operation"];
    mcp::json result;
    
    if (op == "add") {
        if (!params.contains("a") || !params["a"].is_number() ||
            !params.contains("b") || !params["b"].is_number()) {
            Logger::error("mcp.server.toolcall:execute", "Missing required parameters: a, b (numbers)");
            throw std::runtime_error("Missing required parameters: a, b (numbers)");
        }
        
        double a = params["a"];
        double b = params["b"];
        Logger::info("mcp.server.toolcall:execute", "Calculating " + std::to_string(a) + " + " + std::to_string(b));
        result = {{"result", a + b}};
    }
    else if (op == "subtract") {
        if (!params.contains("a") || !params["a"].is_number() ||
            !params.contains("b") || !params["b"].is_number()) {
            throw std::runtime_error("Missing required parameters: a, b (numbers)");
        }
        
        double a = params["a"];
        double b = params["b"];
        Logger::info("mcp.server.toolcall:execute", "Calculating " + std::to_string(a) + " - " + std::to_string(b));
        result = {{"result", a - b}};
    }
    else if (op == "multiply") {
        if (!params.contains("a") || !params["a"].is_number() ||
            !params.contains("b") || !params["b"].is_number()) {
            throw std::runtime_error("Missing required parameters: a, b (numbers)");
        }
        
        double a = params["a"];
        double b = params["b"];
        Logger::info("mcp.server.toolcall:execute", "Calculating " + std::to_string(a) + " * " + std::to_string(b));
        result = {{"result", a * b}};
    }
    else if (op == "divide") {
        if (!params.contains("a") || !params["a"].is_number() ||
            !params.contains("b") || !params["b"].is_number()) {
            throw std::runtime_error("Missing required parameters: a, b (numbers)");
        }
        
        double a = params["a"];
        double b = params["b"];
        
        if (b == 0) {
            Logger::error("mcp.server.toolcall:execute", "Division by zero");
            throw std::runtime_error("Division by zero");
        }
        
        Logger::info("mcp.server.toolcall:execute", "Calculating " + std::to_string(a) + " / " + std::to_string(b));
        result = {{"result", a / b}};
    }
    else {
        Logger::error("mcp.server.toolcall:execute", "Unknown operation: " + op);
        throw std::runtime_error("Unknown operation: " + op);
    }
    
    Logger::info("mcp.server.toolcall:result", "calculator result: " + result.dump());
    
    return result;
}

// Example tool handler for text processing
mcp::json text_processor_handler(const mcp::json& params) {
    Logger::info("mcp.server.toolcall:execute", "Executing text_processor tool with params: " + params.dump());
    
    if (!params.contains("text") || !params["text"].is_string()) {
        Logger::error("mcp.server.toolcall:execute", "Missing required parameter: text (string)");
        throw std::runtime_error("Missing required parameter: text (string)");
    }
    
    std::string text = params["text"];
    mcp::json result = {
        {"original_length", text.length()},
        {"original_text", text}
    };
    
    // Process text based on parameters
    if (params.contains("to_uppercase") && params["to_uppercase"].is_boolean() && params["to_uppercase"]) {
        std::string uppercase_text = text;
        std::transform(uppercase_text.begin(), uppercase_text.end(), uppercase_text.begin(), ::toupper);
        result["uppercase"] = uppercase_text;
        Logger::info("mcp.server.toolcall:execute", "Converted text to uppercase");
    }
    
    if (params.contains("to_lowercase") && params["to_lowercase"].is_boolean() && params["to_lowercase"]) {
        std::string lowercase_text = text;
        std::transform(lowercase_text.begin(), lowercase_text.end(), lowercase_text.begin(), ::tolower);
        result["lowercase"] = lowercase_text;
        Logger::info("mcp.server.toolcall:execute", "Converted text to lowercase");
    }
    
    if (params.contains("reverse") && params["reverse"].is_boolean() && params["reverse"]) {
        std::string reversed_text = text;
        std::reverse(reversed_text.begin(), reversed_text.end());
        result["reversed"] = reversed_text;
        Logger::info("mcp.server.toolcall:execute", "Reversed text");
    }
    
    if (params.contains("word_count") && params["word_count"].is_boolean() && params["word_count"]) {
        // Simple word count by counting spaces
        int count = 1;
        for (char c : text) {
            if (c == ' ') {
                count++;
            }
        }
        result["word_count"] = count;
        Logger::info("mcp.server.toolcall:execute", "Counted words: " + std::to_string(count));
    }
    
    Logger::info("mcp.server.toolcall:result", "text_processor result: " + result.dump());
    
    return result;
}

int main() {
    // 创建和配置服务器
    mcp::server server("localhost", 8080);
    server.set_name("MCP Example Server");
    server.set_cors(true);  // Enable CORS for browser clients
    
    Logger::info("mcp.server", "Starting MCP Example Server");
    
    // Create and register tools
    
    // Time tool
    mcp::tool time_tool = mcp::tool_builder("get_time")
        .with_description("Get the current time")
        .build();
    
    server.register_tool(time_tool, get_time_handler);
    Logger::info("mcp.server", "Registered get_time tool");
    
    // Calculator tool
    mcp::tool calc_tool = mcp::tool_builder("calculator")
        .with_description("Perform mathematical calculations")
        .with_string_param("operation", "Operation to perform: add, subtract, multiply, divide")
        .with_number_param("a", "First operand")
        .with_number_param("b", "Second operand")
        .build();
    
    server.register_tool(calc_tool, calculator_handler);
    Logger::info("mcp.server", "Registered calculator tool");
    
    // Text processor tool
    mcp::tool text_tool = mcp::tool_builder("text_processor")
        .with_description("Process and analyze text")
        .with_string_param("text", "Text to process")
        .with_boolean_param("to_uppercase", "Convert to uppercase", false)
        .with_boolean_param("to_lowercase", "Convert to lowercase", false)
        .with_boolean_param("reverse", "Reverse the text", false)
        .with_boolean_param("word_count", "Count words", false)
        .build();
    
    server.register_tool(text_tool, text_processor_handler);
    Logger::info("mcp.server", "Registered text_processor tool");
    
    // Create and register file resource with logging wrapper
    auto file_resource = std::make_shared<mcp::file_resource>("./files");
    auto logging_file_resource = std::make_shared<LoggingResourceWrapper>(file_resource, "/files");
    server.register_resource("/files", logging_file_resource);
    Logger::info("mcp.server", "Registered file resource at /files");
    
    // Create and register workflow resource with logging wrapper
    auto workflow_resource = std::make_shared<mcp::workflow_resource>();
    auto logging_workflow_resource = std::make_shared<LoggingResourceWrapper>(workflow_resource, "/workflows");
    server.register_resource("/workflows", logging_workflow_resource);
    Logger::info("mcp.server", "Registered workflow resource at /workflows");
    
    // Create and register agent resource with logging wrapper
    auto agent_resource = std::make_shared<mcp::agent_resource>();
    auto logging_agent_resource = std::make_shared<LoggingResourceWrapper>(agent_resource, "/agents");
    server.register_resource("/agents", logging_agent_resource);
    Logger::info("mcp.server", "Registered agent resource at /agents");
    
    // Create and register a custom echo agent
    auto echo_agent_ptr = std::make_shared<examples::echo_agent>("echo");
    echo_agent_ptr->initialize({
        {"uppercase", false},
        {"prefix", true},
        {"prefix_text", "Server Echo: "}
    });
    
    agent_resource->register_agent(echo_agent_ptr);
    Logger::info("mcp.server", "Registered echo agent");
    
    // Create a workflow and register it
    mcp::workflow time_text_workflow("time_with_text", "Get time and process text");
    mcp::tool_registry::instance().register_tool(time_tool, get_time_handler);
    mcp::tool_registry::instance().register_tool(text_tool, text_processor_handler);
    time_text_workflow.add_tool_call("get_time");
    time_text_workflow.add_tool_call("text_processor", {
        {"text", "Current time is important!"},
        {"to_uppercase", true},
        {"word_count", true}
    });
    
    workflow_resource->register_workflow(time_text_workflow);
    Logger::info("mcp.server", "Registered time_with_text workflow");
    
    // Create a workflow agent
    auto workflow_agent_ptr = std::make_shared<examples::workflow_agent>("workflow_runner");
    mcp::tool_registry::instance().register_tool(calc_tool, calculator_handler);
    workflow_agent_ptr->initialize({
        {"workflows", mcp::json::array({
            {
                {"name", "time_text_calc"},
                {"description", "Get time, process text and calculate"},
                {"steps", mcp::json::array({
                    {{"tool_name", "get_time"}},
                    {{"tool_name", "text_processor"}, {"parameters", {
                        {"text", "Let's do some math!"},
                        {"to_uppercase", true}
                    }}},
                    {{"tool_name", "calculator"}, {"parameters", {
                        {"operation", "add"},
                        {"a", 10},
                        {"b", 20}
                    }}}
                })}
            }
        })}
    });
    
    agent_resource->register_agent(workflow_agent_ptr);
    Logger::info("mcp.server", "Registered workflow_runner agent");
    
    // Start the server (blocking)
    std::cout << "Starting server at http://localhost:8080" << std::endl;
    std::cout << "Press Ctrl+C to stop" << std::endl;
    
    try {
        Logger::info("mcp.server", "Server starting...");
        server.start(true);
    } catch (const std::exception& e) {
        Logger::error("mcp.server", std::string("Error: ") + e.what());
        return 1;
    }
    
    return 0;
}