#!/usr/bin/env python3
"""
Test MCP server connection with proper Content-Length headers
"""

import subprocess
import time
import json
import sys
import os
import threading
import queue
import tempfile

class MCPServerTester:
    def __init__(self, server_path):
        self.server_path = server_path
        self.server_process = None
        self.message_queue = queue.Queue()
        
    def start_server(self):
        """Start MCP server process"""
        print(f"Starting MCP server: {self.server_path}")
        self.server_process = subprocess.Popen(
            [self.server_path],
            stdin=subprocess.PIPE,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            text=False,  # Use binary mode for proper handling
            bufsize=0
        )
        
        # Start thread to read stderr
        stderr_thread = threading.Thread(target=self._read_stderr)
        stderr_thread.daemon = True
        stderr_thread.start()
        
        # Start thread to read stdout
        stdout_thread = threading.Thread(target=self._read_stdout)
        stdout_thread.daemon = True
        stdout_thread.start()
        
        time.sleep(1)  # Give server time to start
        
    def _read_stderr(self):
        """Read stderr output from server"""
        while self.server_process and self.server_process.poll() is None:
            try:
                line = self.server_process.stderr.readline()
                if line:
                    print(f"SERVER STDERR: {line.decode().strip()}")
            except:
                break
                
    def _read_stdout(self):
        """Read stdout output from server without Content-Length parsing"""
        buffer = b""
        while self.server_process and self.server_process.poll() is None:
            try:
                data = self.server_process.stdout.readline()
                if not data:
                    break
                buffer += data
                
                # Try to parse JSON directly from each line
                try:
                    message = buffer.strip()
                    if message:
                        parsed = json.loads(message.decode())
                        self.message_queue.put(parsed)
                        print(f"SERVER MESSAGE: {json.dumps(parsed, indent=2)}")
                        buffer = b""
                except json.JSONDecodeError:
                    # If not valid JSON yet, continue reading
                    continue
            except:
                break
                        
    def send_jsonrpc_request(self, method, params=None, id=1):
        """Send a JSON-RPC request without Content-Length header"""
        request = {
            "jsonrpc": "2.0",
            "method": method,
            "id": id,
            "params": {}
        }
        if params:
            request["params"] = params
            
        request_str = json.dumps(request)
        message = f"{request_str}\n"
        print(f"SENDING: {repr(message)}")
        
        self.server_process.stdin.write(message.encode())
        self.server_process.stdin.flush()
        
        # Wait for response with timeout
        try:
            response = self.message_queue.get(timeout=5)
            return response
        except queue.Empty:
            print("TIMEOUT: No response received")
            return None
            
    def test_initialize(self):
        """Test initialize request"""
        print("\n=== Testing Initialize ===")
        
        params = {
            "protocolVersion": "2025-11-25",
            "capabilities": {},
            "clientInfo": {
                "name": "test-client",
                "version": "1.0.0"
            }
        }
        
        print("\nSending initialize request...")
        init_response = self.send_jsonrpc_request("initialize", params, id=1)
        if init_response:
            print(f"Initialize response: {json.dumps(init_response, indent=2)}")
            print("✓ Initialize successful")
            return True
        else:
            print("✗ Initialize failed")
            return False
            
    def test_simple_ping(self):
        """Test simple ping"""
        print("\n=== Testing Simple Ping ===")
        
        # Try a ping
        print("\nSending ping request...")
        ping_response = self.send_jsonrpc_request("ping", id=2)
        if ping_response:
            print(f"Ping response: {json.dumps(ping_response, indent=2)}")
            print("✓ Ping successful")
            return True
        else:
            print("✗ Ping failed")
            return False
        
    def stop_server(self):
        """Stop MCP server process"""
        if self.server_process:
            print("\nStopping server...")
            self.server_process.terminate()
            try:
                self.server_process.wait(timeout=5)
            except subprocess.TimeoutExpired:
                self.server_process.kill()
            self.server_process = None

def main():
    if len(sys.argv) != 2:
        print("Usage: python3 test_content_length.py <server-executable>")
        sys.exit(1)
        
    server_path = sys.argv[1]
    if not os.path.exists(server_path):
        print(f"Error: Server executable '{server_path}' not found")
        sys.exit(1)
        
    tester = MCPServerTester(server_path)
    
    try:
        tester.start_server()
        
        # Wait a bit for server to fully start
        time.sleep(2)
        
        # Test initialize first
        init_success = tester.test_initialize()
        
        if init_success:
            # Test simple ping
            ping_success = tester.test_simple_ping()
            
            if ping_success:
                print("\n✓ Server test completed successfully")
            else:
                print("\n✗ Ping test failed")
        else:
            print("\n✗ Initialize test failed")
            
    except KeyboardInterrupt:
        print("\nTest interrupted by user")
    except Exception as e:
        print(f"\nError during test: {e}")
        import traceback
        traceback.print_exc()
    finally:
        tester.stop_server()

if __name__ == "__main__":
    main()