"""Open one supplied lab server using the pinned Python SDK's stdio client."""
from contextlib import asynccontextmanager
import os
from pathlib import Path
import sys
from mcp import ClientSession, StdioServerParameters
from mcp.client.stdio import stdio_client


@asynccontextmanager
async def lab_session(filename: str = "label_server.py"):
    if filename not in {"label_server.py", "validation_server.py", "context_server.py"}:
        raise ValueError("Choose a supplied lab server")
    server = Path(__file__).resolve().parent / filename
    if not server.is_file():
        raise FileNotFoundError("Save " + filename + " beside lab_session.py")
    parameters = StdioServerParameters(
        command=sys.executable,
        args=["-u", str(server)],
        env={"PYTHONPATH": os.environ.get("PYTHONPATH", "")},
    )
    async with stdio_client(parameters) as (reader, writer):
        async with ClientSession(reader, writer) as session:
            initialized = await session.initialize()
            yield session, initialized