#!/usr/bin/env python3
"""Simple test script for MCP server via stdio"""
import asyncio
import json
import os
import sys
import subprocess
from typing import Any

# Set demo mode
os.environ['DAGSTER_DEMO_MODE'] = 'true'

# Add src to path
sys.path.insert(0, os.path.join(os.path.dirname(__file__), 'src'))

async def test_mcp_tools():
    """Test MCP tools via stdio"""
    # Start the MCP server
    proc = await asyncio.create_subprocess_exec(
        sys.executable,
        'src/mcp_dagster/__main__.py',
        stdin=subprocess.PIPE,
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE,
        env={**os.environ, 'PYTHONPATH': os.path.join(os.path.dirname(__file__), 'src')}
    )
    
    async def send_request(request: dict[str, Any]) -> dict[str, Any]:
        """Send a JSON-RPC request"""
        proc.stdin.write((json.dumps(request) + '\n').encode())
        await proc.stdin.drain()
        
        response_line = await proc.stdout.readline()
        if not response_line:
            raise Exception("No response from server")
        
        return json.loads(response_line.decode())
    
    try:
        # Initialize
        init_req = {
            "jsonrpc": "2.0",
            "id": 1,
            "method": "initialize",
            "params": {
                "protocolVersion": "2024-11-05",
                "capabilities": {},
                "clientInfo": {"name": "test-client", "version": "1.0.0"}
            }
        }
        
        init_response = await send_request(init_req)
        print(f"Initialize response: {init_response.get('result', {}).get('serverInfo', {})}")
        
        # Send initialized notification
        await send_request({
            "jsonrpc": "2.0",
            "method": "notifications/initialized"
        })
        
        # List tools
        tools_req = {
            "jsonrpc": "2.0",
            "id": 2,
            "method": "tools/list"
        }
        
        tools_response = await send_request(tools_req)
        tools = tools_response.get('result', {}).get('tools', [])
        
        print(f"\n✅ Found {len(tools)} tools:")
        for tool in tools:
            print(f"  - {tool.get('name')}: {tool.get('description', 'No description')}")
        
        return len(tools)
        
    finally:
        proc.terminate()
        await proc.wait()


if __name__ == "__main__":
    try:
        count = asyncio.run(test_mcp_tools())
        print(f"\n✅ MCP test passed! Found {count} tools.")
        sys.exit(0)
    except Exception as e:
        print(f"\n❌ MCP test failed: {e}")
        import traceback
        traceback.print_exc()
        sys.exit(1)
