"""
Test script for MCP Server with SSE transport
Tests all endpoints and MCP protocol compliance
"""

import requests
import json
import time

# Configuration
BASE_URL = "http://localhost:7860"  # Change to your HF Space URL for remote testing

def test_health_check():
    """Test health check endpoint"""
    print("🔍 Testing health check...")
    response = requests.get(f"{BASE_URL}/health")
    print(f"Status: {response.status_code}")
    print(f"Response: {json.dumps(response.json(), indent=2)}")
    assert response.status_code == 200
    assert response.json()["status"] == "healthy"
    print("✅ Health check passed\n")


def test_tools_list():
    """Test tools listing endpoint"""
    print("🔍 Testing tools list...")
    response = requests.get(f"{BASE_URL}/tools")
    print(f"Status: {response.status_code}")
    data = response.json()
    print(f"Number of tools: {len(data['tools'])}")
    for tool in data['tools']:
        print(f"  - {tool['name']}")
    assert response.status_code == 200
    assert len(data['tools']) == 7
    print("✅ Tools list passed\n")


def test_mcp_initialize():
    """Test MCP initialize method"""
    print("🔍 Testing MCP initialize...")
    payload = {
        "jsonrpc": "2.0",
        "id": 1,
        "method": "initialize",
        "params": {
            "protocolVersion": "2024-11-05",
            "capabilities": {},
            "clientInfo": {
                "name": "test-client",
                "version": "1.0.0"
            }
        }
    }
    response = requests.post(f"{BASE_URL}/message", json=payload)
    print(f"Status: {response.status_code}")
    print(f"Response: {json.dumps(response.json(), indent=2)}")
    assert response.status_code == 200
    assert response.json()["result"]["serverInfo"]["name"] == "sec-financial-data"
    print("✅ MCP initialize passed\n")


def test_mcp_tools_list():
    """Test MCP tools/list method"""
    print("🔍 Testing MCP tools/list...")
    payload = {
        "jsonrpc": "2.0",
        "id": 2,
        "method": "tools/list"
    }
    response = requests.post(f"{BASE_URL}/message", json=payload)
    print(f"Status: {response.status_code}")
    data = response.json()
    print(f"Number of tools: {len(data['result']['tools'])}")
    assert response.status_code == 200
    assert len(data['result']['tools']) == 7
    print("✅ MCP tools/list passed\n")


def test_search_company():
    """Test search_company tool"""
    print("🔍 Testing search_company tool...")
    payload = {
        "jsonrpc": "2.0",
        "id": 3,
        "method": "tools/call",
        "params": {
            "name": "search_company",
            "arguments": {
                "company_name": "Microsoft"
            }
        }
    }
    response = requests.post(f"{BASE_URL}/message", json=payload)
    print(f"Status: {response.status_code}")
    data = response.json()
    print(f"Response preview: {data['result']['content'][0]['text'][:200]}...")
    assert response.status_code == 200
    assert "content" in data["result"]
    print("✅ search_company passed\n")


def test_get_company_info():
    """Test get_company_info tool"""
    print("🔍 Testing get_company_info tool...")
    payload = {
        "jsonrpc": "2.0",
        "id": 4,
        "method": "tools/call",
        "params": {
            "name": "get_company_info",
            "arguments": {
                "cik": "0000789019"  # Microsoft
            }
        }
    }
    response = requests.post(f"{BASE_URL}/message", json=payload)
    print(f"Status: {response.status_code}")
    data = response.json()
    print(f"Response preview: {data['result']['content'][0]['text'][:200]}...")
    assert response.status_code == 200
    print("✅ get_company_info passed\n")


def test_get_latest_financial_data():
    """Test get_latest_financial_data tool"""
    print("🔍 Testing get_latest_financial_data tool...")
    payload = {
        "jsonrpc": "2.0",
        "id": 5,
        "method": "tools/call",
        "params": {
            "name": "get_latest_financial_data",
            "arguments": {
                "cik": "0000789019"  # Microsoft
            }
        }
    }
    response = requests.post(f"{BASE_URL}/message", json=payload)
    print(f"Status: {response.status_code}")
    data = response.json()
    print(f"Response preview: {data['result']['content'][0]['text'][:300]}...")
    assert response.status_code == 200
    print("✅ get_latest_financial_data passed\n")


def test_extract_financial_metrics():
    """Test extract_financial_metrics tool"""
    print("🔍 Testing extract_financial_metrics tool...")
    payload = {
        "jsonrpc": "2.0",
        "id": 6,
        "method": "tools/call",
        "params": {
            "name": "extract_financial_metrics",
            "arguments": {
                "cik": "0000789019",  # Microsoft
                "years": 2
            }
        }
    }
    response = requests.post(f"{BASE_URL}/message", json=payload)
    print(f"Status: {response.status_code}")
    data = response.json()
    print(f"Response preview: {data['result']['content'][0]['text'][:300]}...")
    assert response.status_code == 200
    print("✅ extract_financial_metrics passed\n")


def test_sse_endpoint():
    """Test SSE endpoint (basic connection)"""
    print("🔍 Testing SSE endpoint...")
    try:
        response = requests.get(f"{BASE_URL}/sse", stream=True, timeout=5)
        print(f"Status: {response.status_code}")
        print("Reading first event...")
        
        # Read first event
        for line in response.iter_lines():
            if line:
                decoded_line = line.decode('utf-8')
                print(f"Received: {decoded_line}")
                if decoded_line.startswith('data:'):
                    break
        
        print("✅ SSE endpoint passed\n")
    except requests.exceptions.Timeout:
        print("⚠️  SSE endpoint timeout (expected for streaming)\n")


def main():
    """Run all tests"""
    print("=" * 60)
    print("🧪 MCP Server SSE Transport Test Suite")
    print("=" * 60)
    print()
    
    tests = [
        ("Health Check", test_health_check),
        ("Tools List", test_tools_list),
        ("MCP Initialize", test_mcp_initialize),
        ("MCP Tools List", test_mcp_tools_list),
        ("Search Company", test_search_company),
        ("Get Company Info", test_get_company_info),
        ("Get Latest Financial Data", test_get_latest_financial_data),
        ("Extract Financial Metrics", test_extract_financial_metrics),
        ("SSE Endpoint", test_sse_endpoint),
    ]
    
    passed = 0
    failed = 0
    
    for name, test_func in tests:
        try:
            test_func()
            passed += 1
        except Exception as e:
            print(f"❌ {name} failed: {e}\n")
            failed += 1
    
    print("=" * 60)
    print(f"✅ Tests passed: {passed}/{len(tests)}")
    if failed > 0:
        print(f"❌ Tests failed: {failed}/{len(tests)}")
    print("=" * 60)


if __name__ == "__main__":
    print("\n⚠️  Make sure the server is running:")
    print("   python mcp_server_sse.py\n")
    time.sleep(1)
    main()