"""Measure CLI/stdio discovery and optional real generation without a resident server."""

from __future__ import annotations

import argparse
import asyncio
import base64
import hashlib
import io
import json
import subprocess
import sys
import time
from pathlib import Path

from mcp import ClientSession
from mcp.client.stdio import StdioServerParameters, stdio_client
from PIL import Image


def gpu_snapshot() -> dict[str, object]:
    """Observe externally; do not initialize CUDA in the client or server."""
    result: dict[str, object] = {}
    for name, query in (
        ("gpu", "--query-gpu=index,memory.used,memory.total"),
        ("processes", "--query-compute-apps=pid,used_gpu_memory"),
    ):
        try:
            completed = subprocess.run(
                ["nvidia-smi", query, "--format=csv,noheader,nounits"],
                capture_output=True,
                text=True,
                timeout=10,
                check=True,
            )
            result[name] = completed.stdout.strip().splitlines()
        except (OSError, subprocess.SubprocessError) as exc:
            result[name] = {"unavailable": str(exc)}
    return result


async def run(args: argparse.Namespace) -> None:
    args.result.parent.mkdir(parents=True, exist_ok=True)
    log_path = args.result.with_suffix(".stderr.log")
    # Fixed seeds belong to this verification process, never the production CLI/schema.
    bootstrap = (
        "import os,sys; print('smoke_pid='+str(os.getpid()),file=sys.stderr,flush=True); "
        "from mcp_sd15.service import secrets; "
        f"secrets.randbelow=lambda _: {args.seed}; "
        "from mcp_sd15.cli import main; raise SystemExit(main())"
    )
    params = StdioServerParameters(
        command=sys.executable,
        args=["-c", bootstrap, "--config", str(args.config.resolve())],
        cwd=str(Path(__file__).resolve().parents[1]),
    )
    report: dict[str, object] = {
        "config": str(args.config.resolve()),
        "seed": args.seed,
        "gpu_before": await asyncio.to_thread(gpu_snapshot),
        "stderr": str(log_path.resolve()),
    }
    try:
        with log_path.open("w", encoding="utf-8") as errlog:
            started = time.perf_counter()
            async with stdio_client(params, errlog=errlog) as (read, write):
                report["spawn_seconds"] = round(time.perf_counter() - started, 3)
                async with ClientSession(
                    read, write, read_timeout_seconds=args.timeout
                ) as client:
                    initialized = time.perf_counter()
                    await client.initialize()
                    report["initialize_seconds"] = round(
                        time.perf_counter() - initialized, 3
                    )
                    listed = time.perf_counter()
                    tools = (await client.list_tools()).tools
                    report["list_seconds"] = round(time.perf_counter() - listed, 3)
                    report["discovery_seconds"] = round(time.perf_counter() - started, 3)
                    names = {tool.name for tool in tools}
                    if names != {"generate_image", "generate_guided_image", "inpaint_image"}:
                        raise RuntimeError(f"unexpected tools: {names}")
                    report["tools"] = sorted(names)
                    report["gpu_after_list"] = await asyncio.to_thread(gpu_snapshot)
                    report["stderr_after_list"] = log_path.read_text(encoding="utf-8")
                    images: list[dict[str, object]] = []
                    report["images"] = images
                    for index in range(args.images):
                        generated = time.perf_counter()
                        result = await client.call_tool(
                            "generate_image",
                            {
                                "positive_prompt": args.prompt,
                                "negative_prompt": "blurry, low quality",
                                "size": "square_512",
                            },
                        )
                        elapsed = time.perf_counter() - generated
                        if getattr(result, "is_error", getattr(result, "isError", False)):
                            raise RuntimeError(f"generation failed: {result}")
                        if len(result.content) != 1 or result.content[0].type != "image":
                            raise RuntimeError(f"unexpected content: {result}")
                        content = result.content[0]
                        if getattr(content, "mime_type", getattr(content, "mimeType", None)) != (
                            "image/png"
                        ):
                            raise RuntimeError("generation did not return image/png")
                        png = base64.b64decode(content.data, validate=True)
                        destination = args.result.with_name(f"{args.result.stem}-{index + 1}.png")
                        destination.write_bytes(png)
                        with Image.open(io.BytesIO(png)) as image:
                            image.load()
                            if image.format != "PNG" or image.size != (512, 512):
                                raise RuntimeError("generation returned unexpected format or size")
                            images.append(
                                {
                                    "elapsed_seconds": round(elapsed, 3),
                                    "format": image.format,
                                    "dimensions": image.size,
                                    "metadata": image.info,
                                    "pixel_sha256": hashlib.sha256(image.tobytes()).hexdigest(),
                                    "output": str(destination.resolve()),
                                    "gpu": await asyncio.to_thread(gpu_snapshot),
                                }
                            )
                    report["stderr_after_calls"] = log_path.read_text(encoding="utf-8")
        report["status"] = "passed"
    except BaseException as exc:
        report["status"] = "failed"
        report["error"] = repr(exc)
        raise
    finally:
        report["gpu_after_exit"] = await asyncio.to_thread(gpu_snapshot)
        report["stderr_final"] = log_path.read_text(encoding="utf-8")
        args.result.write_text(
            json.dumps(report, ensure_ascii=False, indent=2) + "\n", encoding="utf-8"
        )
        print(f"result={args.result.resolve()}")


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--config", type=Path, required=True)
    parser.add_argument("--result", type=Path, required=True)
    parser.add_argument("--images", type=int, choices=(0, 1, 2), default=0)
    parser.add_argument("--seed", type=int, default=2026091201)
    parser.add_argument("--timeout", type=float, default=300)
    parser.add_argument("--prompt", default="a small red cube on a white table, studio photograph")
    asyncio.run(run(parser.parse_args()))


if __name__ == "__main__":
    main()