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

import json
import subprocess
import sys
import time
import os
from threading import Thread

# Configuration - adjust these as needed
HOLOGRES_ENV = {
    "HOLOGRES_HOST": "localhost",
    "HOLOGRES_PORT": "5432",
    "HOLOGRES_USER": "your_user",
    "HOLOGRES_PASSWORD": "your_password",
    "HOLOGRES_DATABASE": "your_database"
}

class StdioClient:
    def __init__(self, command):
        """Initialize the STDIO client with the command to run the server."""
        self.process = subprocess.Popen(
            command,
            stdin=subprocess.PIPE,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            env={**os.environ, **HOLOGRES_ENV},
            text=True,
            bufsize=1
        )
        self.next_id = 1
        
        # Start a thread to read server output
        self.output_thread = Thread(target=self._read_output)
        self.output_thread.daemon = True
        self.output_thread.start()
    
    def _read_output(self):
        """Read and print server output."""
        for line in self.process.stdout:
            try:
                response = json.loads(line)
                print(f"Received: {json.dumps(response, indent=2)}")
            except json.JSONDecodeError:
                print(f"Server output: {line.strip()}")
    
    def send_message(self, method, params=None):
        """Send a JSON-RPC message to the server."""
        message = {
            "jsonrpc": "2.0",
            "id": self.next_id,
            "method": method
        }
        self.next_id += 1
        
        if params:
            message["params"] = params
        
        message_str = json.dumps(message)
        print(f"Sending: {message_str}")
        
        self.process.stdin.write(message_str + "\n")
        self.process.stdin.flush()
        
        # In a real application, you would wait for and parse the response
        # For this example, we just sleep to allow the response to be printed
        time.sleep(1)
    
    def close(self):
        """Close the connection to the server."""
        if self.process:
            self.process.stdin.close()
            self.process.terminate()
            self.process.wait(timeout=5)

def main():
    """Run the example client."""
    # Command to run the server in STDIO mode
    command = ["python", "-m", "hologres_mcp_server.main", "--transport", "stdio"]
    
    client = None
    try:
        # Start the client
        print("Starting Hologres MCP Server in STDIO mode...")
        client = StdioClient(command)
        
        # Example 1: Initialize the connection
        print("\n=== Example 1: Initialize Connection ===")
        client.send_message("initialize", {
            "clientInfo": {
                "name": "python-stdio-test-client",
                "version": "1.0.0"
            }
        })
        
        # Example 2: List available tools
        print("\n=== Example 2: List Tools ===")
        client.send_message("listTools")
        
        # Example 3: Execute a SQL query
        print("\n=== Example 3: Execute SQL Query ===")
        client.send_message("callTool", {
            "name": "execute_sql",
            "arguments": {
                "query": "SELECT 'Hello, Hologres!' AS greeting"
            }
        })
        
        # Example 4: Read a resource
        print("\n=== Example 4: Read Resource ===")
        client.send_message("readResource", {
            "uri": "hologres:///schemas"
        })
        
        print("\nAll examples completed successfully!")
        
    except Exception as e:
        print(f"Error: {str(e)}")
    finally:
        # Close the client
        if client:
            print("Closing client...")
            client.close()

if __name__ == "__main__":
    main()