#!/usr/bin/env python3
"""P0 MCP SDK interoperability and long-running connection regression test."""

from __future__ import annotations

import argparse
import asyncio
import json
import os
import shutil
import tempfile
from pathlib import Path

from mcp.client.session import ClientSession
from mcp.client.stdio import StdioServerParameters, stdio_client
from mcp.types import CallToolResult


def payload(result: CallToolResult) -> dict:
    if result.isError:
        raise AssertionError(f"tool error: {result.content!r}")
    texts = [item.text for item in result.content if item.type == "text"]
    return json.loads("".join(texts))


async def run(server: Path, iterations: int) -> None:
    data_dir = Path(tempfile.mkdtemp(prefix="tiny-convmem-sdk-p0-"))
    params = StdioServerParameters(
        command=str(server.resolve()),
        args=["--data-dir", str(data_dir)],
        env=dict(os.environ),
    )
    completed = 0
    with tempfile.TemporaryFile(mode="w+", encoding="utf-8") as errlog:
        try:
            async with stdio_client(params, errlog=errlog) as (read_stream, write_stream):
                async with ClientSession(read_stream, write_stream) as session:
                    await session.initialize()
                    listed = await session.list_tools()
                    assert {tool.name for tool in listed.tools} == {
                        "write_message",
                        "read_messages",
                        "search_messages",
                        "delete_message",
                        "rebuild_index",
                        "get_system_health",
                    }
                    for index in range(20):
                        payload(
                            await session.call_tool(
                                "write_message",
                                {
                                    "session_id": "target",
                                    "role": "user",
                                    "content": f"sdk needle target {index}",
                                },
                            )
                        )
                    for index in range(10):
                        payload(
                            await session.call_tool(
                                "write_message",
                                {
                                    "session_id": "other",
                                    "role": "user",
                                    "content": f"sdk needle other {index}",
                                },
                            )
                        )

                    for index in range(iterations):
                        arguments = {"query": "needle"}
                        if index % 2:
                            arguments["session_id"] = "target"
                        results = payload(
                            await session.call_tool("search_messages", arguments)
                        )["results"]
                        assert len(results) == 20
                        if index % 2:
                            assert all(item["session_id"] == "target" for item in results)
                        completed += 1
            errlog.flush()
            errlog.seek(0)
            server_stderr = errlog.read()
            if "AddressSanitizer" in server_stderr or "ERROR:" in server_stderr:
                raise AssertionError(f"server diagnostics reported a failure:\n{server_stderr}")
        except Exception as exc:
            errlog.flush()
            errlog.seek(0)
            raise RuntimeError(
                f"failed after {completed}/{iterations} calls; server stderr:\n{errlog.read()}"
            ) from exc
        finally:
            shutil.rmtree(data_dir, ignore_errors=True)


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--server", type=Path, required=True)
    parser.add_argument("--iterations", type=int, default=1000)
    args = parser.parse_args()
    asyncio.run(run(args.server, args.iterations))
    print(f"Python MCP SDK completed {args.iterations} alternating searches")


if __name__ == "__main__":
    main()