"""
Example test cases for Hologres MCP Server.
This file contains test cases that demonstrate how to use the server's features.
"""

import unittest
import requests
import json
import os
import time
import subprocess
import threading
import signal
import sys

# Configuration
SERVER_URL = "http://localhost:8001"
SSE_ENDPOINT = f"{SERVER_URL}/sse"
MESSAGE_ENDPOINT = f"{SERVER_URL}/sse/messages"
HTTP_STREAM_ENDPOINT = f"{SERVER_URL}/mcp"

# Test environment variables
TEST_ENV = {
    "HOLOGRES_HOST": os.environ.get("HOLOGRES_HOST", "localhost"),
    "HOLOGRES_PORT": os.environ.get("HOLOGRES_PORT", "5432"),
    "HOLOGRES_USER": os.environ.get("HOLOGRES_USER", "test_user"),
    "HOLOGRES_PASSWORD": os.environ.get("HOLOGRES_PASSWORD", "test_password"),
    "HOLOGRES_DATABASE": os.environ.get("HOLOGRES_DATABASE", "test_db"),
    "SERVER_HOST": "127.0.0.1",
    "SERVER_PORT": "8001"
}

class ServerProcess:
    """Helper class to manage the server process."""
    
    def __init__(self, transport="sse"):
        self.transport = transport
        self.process = None
    
    def start(self):
        """Start the server process."""
        cmd = [sys.executable, "-m", "hologres_mcp_server.main", "--transport", self.transport]
        self.process = subprocess.Popen(
            cmd,
            env={**os.environ, **TEST_ENV},
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE
        )
        # Give the server time to start
        time.sleep(2)
    
    def stop(self):
        """Stop the server process."""
        if self.process:
            self.process.terminate()
            try:
                self.process.wait(timeout=5)
            except subprocess.TimeoutExpired:
                self.process.kill()
            self.process = None


class TestHologresMCPServer(unittest.TestCase):
    """Test cases for Hologres MCP Server."""
    
    @classmethod
    def setUpClass(cls):
        """Set up the test environment."""
        cls.server = ServerProcess("sse")
        cls.server.start()
    
    @classmethod
    def tearDownClass(cls):
        """Clean up the test environment."""
        cls.server.stop()
    
    def send_message(self, method, params=None):
        """Send a JSON-RPC message to the server."""
        message = {
            "jsonrpc": "2.0",
            "id": int(time.time() * 1000),
            "method": method
        }
        
        if params:
            message["params"] = params
        
        response = requests.post(MESSAGE_ENDPOINT, json=message)
        self.assertEqual(response.status_code, 200, f"Failed to send message: {response.status_code}")
        
        return response.json()
    
    def test_01_initialize(self):
        """Test initializing the connection."""
        response = self.send_message("initialize", {
            "clientInfo": {
                "name": "test-client",
                "version": "1.0.0"
            }
        })
        
        self.assertIn("result", response, "Initialize response should contain 'result'")
        self.assertIn("protocolVersion", response["result"], "Result should contain 'protocolVersion'")
        self.assertIn("capabilities", response["result"], "Result should contain 'capabilities'")
        self.assertIn("serverInfo", response["result"], "Result should contain 'serverInfo'")
    
    def test_02_list_tools(self):
        """Test listing available tools."""
        response = self.send_message("listTools")
        
        self.assertIn("result", response, "ListTools response should contain 'result'")
        self.assertIn("tools", response["result"], "Result should contain 'tools'")
        
        tools = response["result"]["tools"]
        self.assertIsInstance(tools, list, "Tools should be a list")
        
        # Check for expected tools
        tool_names = [tool["name"] for tool in tools]
        expected_tools = ["execute_sql", "analyze_table", "get_query_plan", "get_execution_plan"]
        for tool in expected_tools:
            self.assertIn(tool, tool_names, f"Tool '{tool}' should be available")
    
    def test_03_read_resource(self):
        """Test reading a resource."""
        response = self.send_message("readResource", {
            "uri": "hologres:///schemas"
        })
        
        self.assertIn("result", response, "ReadResource response should contain 'result'")
        self.assertIn("content", response["result"], "Result should contain 'content'")
    
    def test_04_call_tool(self):
        """Test calling a tool."""
        # This is a simple test that should work even without a real database connection
        response = self.send_message("callTool", {
            "name": "execute_sql",
            "arguments": {
                "query": "SELECT 'Hello, Hologres!' AS greeting"
            }
        })
        
        self.assertIn("result", response, "CallTool response should contain 'result'")
        self.assertIn("output", response["result"], "Result should contain 'output'")


class TestHologresMCPServerSTDIO(unittest.TestCase):
    """Test cases for Hologres MCP Server in STDIO mode."""
    
    def test_stdio_mode(self):
        """Test that the server can run in STDIO mode."""
        # This is a simplified test that just checks if the server starts in STDIO mode
        server = ServerProcess("stdio")
        try:
            server.start()
            # If we got here without an exception, the server started successfully
            self.assertTrue(True, "Server started in STDIO mode")
        finally:
            server.stop()


class TestHologresMCPServerHTTPStream(unittest.TestCase):
    """Test cases for Hologres MCP Server in HTTP Stream mode."""
    
    @classmethod
    def setUpClass(cls):
        """Set up the test environment."""
        cls.server = ServerProcess("http-stream")
        cls.server.start()
    
    @classmethod
    def tearDownClass(cls):
        """Clean up the test environment."""
        cls.server.stop()
    
    def send_batch_request(self, requests_list):
        """Send a batch of JSON-RPC requests to the server."""
        response = requests.post(HTTP_STREAM_ENDPOINT, json=requests_list)
        self.assertEqual(response.status_code, 200, f"Failed to send batch request: {response.status_code}")
        return response.json()
    
    def send_message(self, method, params=None):
        """Send a single JSON-RPC message to the server."""
        message = {
            "jsonrpc": "2.0",
            "id": int(time.time() * 1000),
            "method": method
        }
        
        if params:
            message["params"] = params
        
        response = requests.post(HTTP_STREAM_ENDPOINT, json=message)
        self.assertEqual(response.status_code, 200, f"Failed to send message: {response.status_code}")
        
        return response.json()
    
    def test_01_initialize(self):
        """Test initializing the connection."""
        response = self.send_message("initialize", {
            "clientInfo": {
                "name": "test-client",
                "version": "1.0.0"
            }
        })
        
        self.assertIn("result", response, "Initialize response should contain 'result'")
        self.assertIn("protocolVersion", response["result"], "Result should contain 'protocolVersion'")
        self.assertIn("capabilities", response["result"], "Result should contain 'capabilities'")
        self.assertIn("serverInfo", response["result"], "Result should contain 'serverInfo'")
    
    def test_02_list_tools(self):
        """Test listing available tools."""
        response = self.send_message("listTools")
        
        self.assertIn("result", response, "ListTools response should contain 'result'")
        self.assertIn("tools", response["result"], "Result should contain 'tools'")
        
        tools = response["result"]["tools"]
        self.assertIsInstance(tools, list, "Tools should be a list")
        
        # Check for expected tools
        tool_names = [tool["name"] for tool in tools]
        expected_tools = ["execute_sql", "analyze_table", "get_query_plan", "get_execution_plan"]
        for tool in expected_tools:
            self.assertIn(tool, tool_names, f"Tool '{tool}' should be available")
    
    def test_03_read_resource(self):
        """Test reading a resource."""
        response = self.send_message("readResource", {
            "uri": "hologres:///schemas"
        })
        
        self.assertIn("result", response, "ReadResource response should contain 'result'")
        self.assertIn("content", response["result"], "Result should contain 'content'")
    
    def test_04_call_tool(self):
        """Test calling a tool."""
        # This is a simple test that should work even without a real database connection
        response = self.send_message("callTool", {
            "name": "execute_sql",
            "arguments": {
                "query": "SELECT 'Hello, Hologres!' AS greeting"
            }
        })
        
        self.assertIn("result", response, "CallTool response should contain 'result'")
        self.assertIn("output", response["result"], "Result should contain 'output'")
    
    def test_05_batch_request(self):
        """Test batch request functionality."""
        batch_requests = [
            {
                "jsonrpc": "2.0",
                "id": 1,
                "method": "initialize",
                "params": {
                    "clientInfo": {
                        "name": "test-client",
                        "version": "1.0.0"
                    }
                }
            },
            {
                "jsonrpc": "2.0",
                "id": 2,
                "method": "listTools"
            },
            {
                "jsonrpc": "2.0",
                "id": 3,
                "method": "callTool",
                "params": {
                    "name": "execute_sql",
                    "arguments": {
                        "query": "SELECT 'Hello, Hologres!' AS greeting"
                    }
                }
            }
        ]
        
        responses = self.send_batch_request(batch_requests)
        
        self.assertEqual(len(responses), 3, "Should receive 3 responses")
        self.assertEqual(responses[0]["id"], 1, "First response should have id 1")
        self.assertEqual(responses[1]["id"], 2, "Second response should have id 2")
        self.assertEqual(responses[2]["id"], 3, "Third response should have id 3")
        
        self.assertIn("result", responses[0], "First response should contain 'result'")
        self.assertIn("result", responses[1], "Second response should contain 'result'")
        self.assertIn("result", responses[2], "Third response should contain 'result'")


if __name__ == "__main__":
    unittest.main()