#!/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

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=True,
            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:
            line = self.server_process.stderr.readline()
            if line:
                print(f"SERVER STDERR: {line.strip()}")
                
    def _read_stdout(self):
        """Read stdout output from server without Content-Length parsing"""
        while self.server_process and self.server_process.poll() is None:
            try:
                line = self.server_process.stdout.readline()
                if not line:
                    break
                
                # Try to parse JSON directly from each line
                try:
                    message = line.strip()
                    if message:
                        parsed = json.loads(message)
                        self.message_queue.put(parsed)
                        print(f"SERVER MESSAGE: {json.dumps(parsed, indent=2)}")
                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)
        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_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=1)
        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_mcp_with_headers.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 simple ping
        success = tester.test_simple_ping()
        
        if success:
            print("\n✓ Server test completed successfully")
        else:
            print("\n✗ Server 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()