"""
Example Python client for testing the Hologres MCP Server in SSE mode.
This script demonstrates how to connect to the server and use its features.
"""

import requests
import json
import time
import os
import sseclient

# Configuration
SERVER_URL = "http://localhost:8001"  # Change to your server's URL
SSE_ENDPOINT = f"{SERVER_URL}/sse"
MESSAGE_ENDPOINT = None  # Will be set after connecting to SSE endpoint

def connect_to_sse():
    """Connect to the SSE endpoint and get the message endpoint."""
    global MESSAGE_ENDPOINT
    
    print(f"Connecting to SSE endpoint: {SSE_ENDPOINT}")
    response = requests.get(SSE_ENDPOINT, stream=True)
    
    if response.status_code != 200:
        raise Exception(f"Failed to connect to SSE endpoint: {response.status_code}")
    
    client = sseclient.SSEClient(response)
    
    # Get the message endpoint from the first event
    for event in client.events():
        if event.event == "endpoint":
            MESSAGE_ENDPOINT = f"{SERVER_URL}{event.data}"
            print(f"Received message endpoint: {MESSAGE_ENDPOINT}")
            break
    
    if not MESSAGE_ENDPOINT:
        raise Exception("Failed to receive message endpoint")
    
    return client

def send_message(method, params=None):
    """Send a JSON-RPC message to the server."""
    if not MESSAGE_ENDPOINT:
        raise Exception("Not connected to server. Call connect_to_sse() first.")
    
    message = {
        "jsonrpc": "2.0",
        "id": int(time.time() * 1000),
        "method": method
    }
    
    if params:
        message["params"] = params
    
    print(f"Sending message: {json.dumps(message, indent=2)}")
    response = requests.post(MESSAGE_ENDPOINT, json=message)
    
    if response.status_code != 200:
        raise Exception(f"Failed to send message: {response.status_code}")
    
    result = response.json()
    print(f"Received response: {json.dumps(result, indent=2)}")
    return result

def initialize():
    """Initialize the connection to the server."""
    return send_message("initialize", {
        "clientInfo": {
            "name": "python-test-client",
            "version": "1.0.0"
        }
    })

def list_tools():
    """List available tools."""
    return send_message("listTools")

def call_tool(name, arguments):
    """Call a tool with arguments."""
    return send_message("callTool", {
        "name": name,
        "arguments": arguments
    })

def read_resource(uri):
    """Read a resource by URI."""
    return send_message("readResource", {
        "uri": uri
    })

def main():
    """Run the example client."""
    try:
        # Connect to SSE endpoint
        sse_client = connect_to_sse()
        
        # Initialize the connection
        initialize()
        
        # List available tools
        tools_response = list_tools()
        
        # Example 1: Execute a SQL query
        print("\n=== Example 1: Execute SQL Query ===")
        sql_result = call_tool("execute_sql", {
            "query": "SELECT 'Hello, Hologres!' AS greeting"
        })
        
        # Example 2: List all schemas
        print("\n=== Example 2: List All Schemas ===")
        schemas_result = read_resource("hologres:///schemas")
        
        # Example 3: Get query plan
        print("\n=== Example 3: Get Query Plan ===")
        plan_result = call_tool("get_query_plan", {
            "query": "SELECT * FROM information_schema.tables LIMIT 10"
        })
        
        print("\nAll examples completed successfully!")
        
    except Exception as e:
        print(f"Error: {str(e)}")
    finally:
        # Keep the connection open to receive events (in a real application)
        # In this example, we'll just exit
        print("Exiting...")

if __name__ == "__main__":
    main()