#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
Test SSE Connection to Maya MCP Server
This script tests the SSE connection to the Maya MCP server.
Run this script outside of Maya to avoid thread blocking issues.

Usage:
    python test_sse_connection.py
"""

import http.client
import time
import sys
import socket
import json

def test_sse_connection(host="127.0.0.1", port=4550, path="/", timeout=30):
    """
    Test SSE connection to Maya MCP server
    
    Args:
        host (str): Server host
        port (int): Server port
        path (str): Server path
        timeout (int): Connection timeout in seconds
    
    Returns:
        bool: Whether connection was successful
    """
    print(f"Testing SSE connection to {host}:{port}{path}")
    print("-" * 50)
    
    try:
        # Create connection with timeout
        print(f"Creating connection with timeout of {timeout} seconds...")
        conn = http.client.HTTPConnection(host, port, timeout=timeout)
        
        # Send request with SSE headers
        print(f"Sending request to {host}:{port}{path}")
        conn.request("GET", path, headers={"Accept": "text/event-stream"})
        
        # Get response
        print("Waiting for response...")
        response = conn.getresponse()
        print(f"Connection status: {response.status} {response.reason}")
        
        if response.status != 200:
            print(f"Error: Received status code {response.status}")
            return False
        
        # Read initial data
        print("Reading initial data (waiting up to 10 seconds)...")
        start_time = time.time()
        initial_data = ""
        received_complete_data = False
        event_count = 0
        max_events = 5  # Read 5 events
        
        # Try to read data for up to 10 seconds
        while time.time() - start_time < 10 and not received_complete_data and event_count < max_events:
            try:
                # Read data in smaller chunks to avoid truncation
                chunk = response.read(128).decode('utf-8')
                if chunk:
                    initial_data += chunk
                    print(f"Received chunk: {chunk}")
                    
                    # Count events
                    if "event:" in chunk:
                        event_count += 1
                    
                    # Check if we've received complete JSON data
                    if '}' in chunk and 'data:' in initial_data and event_count >= max_events:
                        received_complete_data = True
                        break
                
                # Reduce sleep time to improve responsiveness
                time.sleep(0.1)
                # Periodically print waiting information
                if (time.time() - start_time) % 1 < 0.1:
                    print(f"Waiting for data... ({int(time.time() - start_time)} seconds elapsed)")
            except socket.timeout:
                print("Socket timeout while reading initial data, retrying...")
                continue
        
        # Print complete received data
        if initial_data:
            print("\nComplete received data:")
            print("-" * 50)
            print(initial_data)
            print("-" * 50)
            
            # Try to parse events
            try:
                events = []
                current_event = {}
                for line in initial_data.split('\n'):
                    line = line.strip()
                    if not line or line.startswith(':'):
                        continue
                    
                    if line.startswith('event:'):
                        if current_event and 'event' in current_event:
                            events.append(current_event)
                            current_event = {}
                        current_event['event'] = line[6:].strip()
                    elif line.startswith('data:'):
                        if 'data' not in current_event:
                            current_event['data'] = line[5:].strip()
                        else:
                            current_event['data'] += line[5:].strip()
                
                if current_event and 'event' in current_event:
                    events.append(current_event)
                
                print("\nParsed Events:")
                print("-" * 50)
                for i, event in enumerate(events):
                    print(f"Event {i+1}: {event['event']}")
                    try:
                        data = json.loads(event['data'])
                        print(f"Data: {json.dumps(data, indent=2)}")
                    except:
                        print(f"Data: {event['data']}")
                    print()
                print("-" * 50)
            except Exception as e:
                print(f"Error parsing events: {e}")
        else:
            print("Warning: No initial data received after 10 seconds")
            print("This doesn't necessarily mean the connection failed.")
            print("The server might be configured to not send initial data.")
            print("Connection is established (200 OK), which is a good sign.")
        
        # Close connection
        print("\nClosing connection...")
        conn.close()
        print("Connection closed")
        
        return True
    except socket.timeout as e:
        print(f"Socket timeout error: {e}")
        return False
    except ConnectionRefusedError as e:
        print(f"Connection refused: {e}")
        print("Make sure the server is running and accessible")
        return False
    except Exception as e:
        print(f"Error: {e}")
        import traceback
        traceback.print_exc()
        return False

if __name__ == "__main__":
    # Get command line arguments
    host = "127.0.0.1"
    port = 4550
    path = "/"
    
    # Parse command line arguments
    if len(sys.argv) > 1:
        host = sys.argv[1]
    if len(sys.argv) > 2:
        port = int(sys.argv[2])
    if len(sys.argv) > 3:
        path = sys.argv[3]
    
    # Test connection
    success = test_sse_connection(host, port, path)
    
    # Print result
    print("\n" + "=" * 50)
    if success:
        print("SSE connection test SUCCESSFUL")
    else:
        print("SSE connection test FAILED")
    print("=" * 50)