#!/usr/bin/env python3
"""
Test script to verify MCP server handles concurrent agent requests correctly.
"""
import json
import subprocess
import time
from threading import Thread
from queue import Queue

SILVERSIGHT = "/home/allaun/SilverSight"

def send_mcp_request(request):
    """Send a request to the MCP server and get response."""
    proc = subprocess.Popen(
        ["python3", f"{SILVERSIGHT}/scripts/mcp_autoproof.py", "--stdio"],
        stdin=subprocess.PIPE,
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE,
        text=True,
        cwd=SILVERSIGHT
    )
    
    # Send request
    proc.stdin.write(json.dumps(request) + "\n")
    proc.stdin.flush()
    proc.stdin.close()
    
    # Get response
    response = proc.stdout.readline()
    proc.wait()
    
    return json.loads(response) if response else None

def worker(agent_id, request, results):
    """Worker function for concurrent testing."""
    print(f"[Agent {agent_id}] Sending request: {request.get('method')}")
    start = time.time()
    
    try:
        response = send_mcp_request(request)
        elapsed = time.time() - start
        
        if response:
            result = response.get("result", {})
            status = result.get("content", [{}])[0].get("text", "")
            status_json = json.loads(status) if status else {}
            
            results.put({
                "agent_id": agent_id,
                "elapsed": elapsed,
                "status": status_json.get("status", "unknown"),
                "message": status_json.get("message", "")[:100]
            })
            print(f"[Agent {agent_id}] Response in {elapsed:.2f}s: {status_json.get('status')}")
        else:
            results.put({
                "agent_id": agent_id,
                "elapsed": elapsed,
                "status": "error",
                "message": "No response"
            })
            print(f"[Agent {agent_id}] No response")
    except Exception as e:
        elapsed = time.time() - start
        results.put({
            "agent_id": agent_id,
            "elapsed": elapsed,
            "status": "error",
            "message": str(e)[:100]
        })
        print(f"[Agent {agent_id}] Error: {e}")

def test_concurrent_get_context():
    """Test multiple agents reading the same file."""
    print("\n=== Test: Concurrent get_sorry_context ===")
    results = Queue()
    threads = []
    
    for i in range(3):
        request = {
            "method": "tools/call",
            "params": {
                "name": "get_sorry_context",
                "arguments": {
                    "filepath": "formal/CoreFormalism/ComplexProjectiveSpace.lean"
                }
            }
        }
        t = Thread(target=worker, args=(i, request, results))
        threads.append(t)
        t.start()
        time.sleep(0.1)  # Small delay to stagger requests
    
    for t in threads:
        t.join()
    
    print("\nResults:")
    while not results.empty():
        r = results.get()
        print(f"  Agent {r['agent_id']}: {r['status']} ({r['elapsed']:.2f}s) - {r['message'][:80]}")

def test_concurrent_check_proof():
    """Test multiple agents checking proofs simultaneously."""
    print("\n=== Test: Concurrent check_proof ===")
    results = Queue()
    threads = []
    
    for i in range(2):
        request = {
            "method": "tools/call",
            "params": {
                "name": "check_proof",
                "arguments": {
                    "target": "CoreFormalism.ComplexProjectiveSpace"
                }
            }
        }
        t = Thread(target=worker, args=(i, request, results))
        threads.append(t)
        t.start()
    
    for t in threads:
        t.join()
    
    print("\nResults:")
    while not results.empty():
        r = results.get()
        print(f"  Agent {r['agent_id']}: {r['status']} ({r['elapsed']:.2f}s) - {r['message'][:80]}")

def test_mixed_operations():
    """Test mixed operations happening concurrently."""
    print("\n=== Test: Mixed Operations ===")
    results = Queue()
    threads = []
    
    requests = [
        {
            "method": "tools/call",
            "params": {
                "name": "get_sorry_context",
                "arguments": {
                    "filepath": "formal/CoreFormalism/ComplexProjectiveSpace.lean"
                }
            }
        },
        {
            "method": "tools/call",
            "params": {
                "name": "check_proof",
                "arguments": {
                    "target": "CoreFormalism.ComplexProjectiveSpace"
                }
            }
        },
        {
            "method": "tools/call",
            "params": {
                "name": "list_tools",
                "arguments": {}
            }
        }
    ]
    
    for i, request in enumerate(requests):
        t = Thread(target=worker, args=(i, request, results))
        threads.append(t)
        t.start()
        time.sleep(0.05)
    
    for t in threads:
        t.join()
    
    print("\nResults:")
    while not results.empty():
        r = results.get()
        print(f"  Agent {r['agent_id']}: {r['status']} ({r['elapsed']:.2f}s) - {r['message'][:80]}")

def main():
    print("MCP Server Concurrency Test Suite")
    print("=" * 50)
    
    # Test 1: Concurrent read operations
    test_concurrent_get_context()
    
    # Test 2: Concurrent build operations (should serialize)
    test_concurrent_check_proof()
    
    # Test 3: Mixed operations
    test_mixed_operations()
    
    print("\n" + "=" * 50)
    print("All tests completed!")

if __name__ == "__main__":
    main()