import subprocess
import time
import os
import asyncio

# 全局常量
AGENT_SCRIPT = "compliance-mcp-agent/agent_main_loop.py"

async def log_subprocess_output(stream, prefix):
    """异步读取并打印子进程的输出流。"""
    while True:
        try:
            line = await stream.readline()
            if line:
                print(f"[{prefix}] {line.decode().strip()}")
            else:
                # End of stream
                break
        except Exception as e:
            print(f"Error reading stream for {prefix}: {e}")
            break

async def start_servers():
    """异步启动所有MCP服务器，并捕获它们的日志。"""
    print("="*20 + " Starting MCP Servers " + "="*20)
    
    server_processes = {}
    log_tasks = []
    
    server_scripts = [
        "servers/APICallerServer.py",
        "servers/SchemaValidatorServer.py",
        "servers/DMSProviderServer.py",
        "servers/TestManagerServer.py",
    ]

    project_root = os.path.dirname(os.path.abspath(__file__))

    for script_path_rel in server_scripts:
        script_path_abs = os.path.join(project_root, script_path_rel)
        server_name = os.path.splitext(os.path.basename(script_path_abs))[0]
        server_dir = os.path.dirname(script_path_abs)
        
        print(f"Starting server: {script_path_rel}...")
        
        env = os.environ.copy()
        
        process = await asyncio.create_subprocess_exec(
            "uv", "run", "python", os.path.basename(script_path_abs),
            stdout=asyncio.subprocess.PIPE,
            stderr=asyncio.subprocess.PIPE,
            env=env,
            cwd=server_dir
        )
        
        server_processes[server_name] = process
        print(f"  -> Started {script_path_rel} with PID: {process.pid}")

        log_tasks.append(asyncio.create_task(log_subprocess_output(process.stdout, server_name)))
        log_tasks.append(asyncio.create_task(log_subprocess_output(process.stderr, f"{server_name}-ERROR")))

    print(f"\nAll {len(server_scripts)} servers are running in the background.")
    return server_processes, log_tasks

async def run_agent(agent_script_path):
    """异步运行Agent主循环并实时打印其输出。"""
    print("\n" + "="*20 + " Running Agent " + "="*20)
    
    process = await asyncio.create_subprocess_exec(
        "uv", "run", "python", agent_script_path,
        stdout=asyncio.subprocess.PIPE,
        stderr=asyncio.subprocess.PIPE
    )
    
    agent_stdout_task = asyncio.create_task(log_subprocess_output(process.stdout, "Agent"))
    agent_stderr_task = asyncio.create_task(log_subprocess_output(process.stderr, "Agent-ERROR"))
    
    await process.wait()
    
    await asyncio.gather(agent_stdout_task, agent_stderr_task)
        
    print(f"\nAgent process finished with return code: {process.returncode}")


def cleanup_servers(processes):
    """停止所有服务器进程。"""
    print("\n" + "="*20 + " Cleaning up Servers " + "="*20)
    for name, process in processes.items():
        if process.returncode is None: # 仅当进程仍在运行时才终止
            print(f"Terminating server {name} (PID: {process.pid})...")
            try:
                process.terminate()
                print(f"  -> Terminated signal sent.")
            except ProcessLookupError:
                print(f"  -> Process {name} (PID: {process.pid}) already gone.")
        else:
            print(f"  -> Server {name} (PID: {process.pid}) already finished with code {process.returncode}.")
            
    print("Cleanup complete.")

async def main():
    """
    主入口点，异步运行所有测试。
    """
    server_processes, log_tasks = await start_servers()
    
    print("\nWaiting for servers to initialize...")
    await asyncio.sleep(8)
    
    try:
        await run_agent(AGENT_SCRIPT)
    finally:
        cleanup_servers(server_processes)
        
        # 取消仍在运行的日志任务
        for task in log_tasks:
            task.cancel()
        await asyncio.gather(*log_tasks, return_exceptions=True)
        print("Log watchers terminated.")

if __name__ == "__main__":
    try:
        asyncio.run(main())
    except KeyboardInterrupt:
        print("\nMain process interrupted by user. Cleaning up...")