/**
 * @file test_mcp_server.cpp
 * @brief Test MCP server related functionality
 * 
 * This file contains unit tests for the MCP server module, based on the 2024-11-05 specification.
 */

#include "mcp_server.h"
#include "mcp_client.h"
#include <gtest/gtest.h>
#include <gmock/gmock.h>
#include <string>
#include <thread>
#include <future>
#include <chrono>

// Mock tool handler function
mcp::json test_tool_handler(const mcp::json& params) {
    if (params.contains("input")) {
        return {
            {
                {"type", "text"},
                {"text", "Result: " + params["input"].get<std::string>()}
            }
        };
    } else {
        return {
            {
                {"type", "text"},
                {"text", "Default result"}
            }
        };
    }
}

// Server test class
class ServerTest : public ::testing::Test {
protected:
    void SetUp() override {
        // Create server
        server = std::make_unique<mcp::server>("localhost", 8095);
        server->set_server_info("TestServer", "2024-11-05");
        
        // Set server capabilities
        mcp::json capabilities = {
            {"tools", {{"listChanged", true}}}
        };
        server->set_capabilities(capabilities);
    }
    
    void TearDown() override {
        // Stop server
        if (server && server_thread.joinable()) {
            server->stop();
            server_thread.join();
        }
    }
    
    // Start server
    void start_server() {
        server_thread = std::thread([this]() {
            server->start(false);
        });
        // Wait for server to start
        std::this_thread::sleep_for(std::chrono::milliseconds(100));
    }
    
    std::unique_ptr<mcp::server> server;
    std::thread server_thread;
};

// Test tool registration and retrieval
TEST_F(ServerTest, ToolRegistrationTest) {
    // Create tool
    mcp::tool test_tool = mcp::tool_builder("test_tool")
        .with_description("Test tool")
        .with_string_param("input", "Input parameter")
        .build();
    
    // Register tool
    server->register_tool(test_tool, test_tool_handler);
    
    // Start server
    start_server();
    
    // Create client to verify tool registration
    mcp::client client("localhost", 8095);
    client.set_timeout(5);
    client.initialize("TestClient", mcp::MCP_VERSION);
    
    // Get tool list
    auto tools = client.get_tools();
    EXPECT_EQ(tools.size(), 1);
    EXPECT_EQ(tools[0].name, "test_tool");
    EXPECT_EQ(tools[0].description, "Test tool");
}

// Test method registration and handling
TEST_F(ServerTest, MethodRegistrationTest) {
    // Register method - register as tool since client doesn't have call_method
    server->register_tool(
        mcp::tool_builder("test_method")
            .with_description("Test method")
            .with_string_param("param1", "Parameter 1", false)
            .build(),
        [](const mcp::json& params) -> mcp::json {
            return {
                {
                    {"type", "text"},
                    {"text", "Method call successful"}
                },
                {
                    {"type", "text"},
                    {"text", "params: " + params.dump()}
                }
            };
        }
    );
    
    // Start server
    start_server();
    
    // Create client to verify method registration
    mcp::client client("localhost", 8095);
    client.set_timeout(5);
    client.initialize("TestClient", mcp::MCP_VERSION);
    
    // Call method (via tool)
    mcp::json result = client.call_tool("test_method", {{"param1", "value1"}});
    EXPECT_EQ(result["content"][0]["text"], "Method call successful");
    EXPECT_EQ(result["content"][1]["text"], "params: {\"param1\":\"value1\"}");
    
    // Call non-existent method
    EXPECT_THROW(client.call_tool("non_existent_method"), mcp::mcp_exception);
}

// Test server-client interaction
TEST_F(ServerTest, ServerClientInteractionTest) {
    // Register tool
    mcp::tool test_tool = mcp::tool_builder("test_tool")
        .with_description("Test tool")
        .with_string_param("input", "Input parameter")
        .build();
    server->register_tool(test_tool, test_tool_handler);
    
    // Start server
    start_server();
    
    // Create client
    mcp::client client("localhost", 8095);
    client.set_timeout(5);
    
    // Initialize client
    bool initialized = client.initialize("TestClient", mcp::MCP_VERSION);
    EXPECT_TRUE(initialized);
    
    // Get tool list
    auto tools = client.get_tools();
    EXPECT_EQ(tools.size(), 1);
    EXPECT_EQ(tools[0].name, "test_tool");
    
    // Call tool
    mcp::json tool_result = client.call_tool("test_tool", {{"input", "Test input"}});
    EXPECT_EQ(tool_result["content"][0]["text"], "Result: Test input");
}

// Test notification functionality
TEST_F(ServerTest, NotificationTest) {
    // Start server
    start_server();
    
    // Create client
    mcp::client client("localhost", 8095);
    client.initialize("TestClient", mcp::MCP_VERSION);
    
    // Add tool on server side, triggering notification
    mcp::tool new_tool = mcp::tool_builder("new_tool")
        .with_description("New tool")
        .build();
    server->register_tool(new_tool, [](const mcp::json& params) { return mcp::json::object(); });
    
    // Wait for notification processing
    std::this_thread::sleep_for(std::chrono::milliseconds(200));
    
    // Verify tools were added
    auto tools = client.get_tools();
    EXPECT_EQ(tools.size(), 1);
    EXPECT_EQ(tools[0].name, "new_tool");
}

// Test error handling
TEST_F(ServerTest, ErrorHandlingTest) {
    // Register a tool handler that throws an exception
    mcp::tool error_tool = mcp::tool_builder("error_tool")
        .with_description("Error tool")
        .build();
    
    server->register_tool(error_tool, [](const mcp::json& params) -> mcp::json {
        throw std::runtime_error("Test exception");
    });
    
    // Start server
    start_server();
    
    // Create client
    mcp::client client("localhost", 8095);
    client.initialize("TestClient", mcp::MCP_VERSION);
    
    // Call tool that throws exception
    mcp::json result = client.call_tool("error_tool");
    EXPECT_EQ(result.contains("isError"), true);
    EXPECT_EQ(result["isError"], true);
}

// Test server stop and restart
TEST_F(ServerTest, ServerRestartTest) {
    // Start server
    start_server();
    
    // Create client and connect
    mcp::client client("localhost", 8095);
    bool initialized = client.initialize("TestClient", mcp::MCP_VERSION);
    EXPECT_TRUE(initialized);
    
    // Stop server
    server->stop();
    server_thread.join();
    
    // Wait to ensure server is fully stopped
    std::this_thread::sleep_for(std::chrono::milliseconds(500));
    
    // Try to connect, should fail - but may not throw exception, might just return false
    mcp::client client2("localhost", 8095);
    client2.set_timeout(1); // Set shorter timeout
    
    // Try to initialize, should fail
    bool init_result = client2.initialize("TestClient", mcp::MCP_VERSION);
    EXPECT_FALSE(init_result);
    
    // Restart server
    start_server();
    
    // Try to connect again, should succeed
    mcp::client client3("localhost", 8095);
    initialized = client3.initialize("TestClient", mcp::MCP_VERSION);
    EXPECT_TRUE(initialized);
}