"""Command-line entry point for safe Brain runtime diagnostics."""

from __future__ import annotations

import argparse
import json
import platform
from collections.abc import Sequence
from typing import Any

from yovision_brain import __version__


def _load_torch() -> Any:
    try:
        import torch
    except ImportError as exc:
        raise RuntimeError(
            "PyTorch runtime is not installed; install the pinned CPU or CUDA wheel "
            "from Brain/README.md"
        ) from exc
    return torch


def runtime_info() -> dict[str, object]:
    """Return non-sensitive interpreter and compute-runtime facts."""
    info: dict[str, object] = {
        "python": platform.python_version(),
        "torch_installed": False,
        "torch_version": None,
        "cuda_build": None,
        "cuda_available": False,
        "cuda_device_count": 0,
    }
    try:
        torch = _load_torch()
    except RuntimeError:
        return info

    cuda_available = bool(torch.cuda.is_available())
    info.update(
        {
            "torch_installed": True,
            "torch_version": torch.__version__,
            "cuda_build": torch.version.cuda,
            "cuda_available": cuda_available,
            "cuda_device_count": torch.cuda.device_count() if cuda_available else 0,
        }
    )
    return info


def smoke(device: str) -> dict[str, object]:
    """Run a deterministic tensor operation on exactly the requested device."""
    torch = _load_torch()
    if device == "cuda" and not torch.cuda.is_available():
        raise RuntimeError("CUDA was requested but PyTorch reports no available CUDA device")

    tensor = torch.tensor([[1.0, 2.0], [3.0, 4.0]], device=device)
    result = tensor @ tensor
    expected = torch.tensor([[7.0, 10.0], [15.0, 22.0]], device=device)
    if not torch.equal(result, expected):
        raise RuntimeError("tensor smoke result did not match the expected value")

    return {
        "status": "ok",
        "device": device,
        "torch_version": torch.__version__,
        "cuda_build": torch.version.cuda,
    }


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        prog="yovision-brain",
        description="YoVision Brain runtime diagnostics",
    )
    parser.add_argument("--version", action="version", version=f"%(prog)s {__version__}")
    action = parser.add_mutually_exclusive_group()
    action.add_argument(
        "--runtime-info",
        action="store_true",
        help="print non-sensitive Python/PyTorch/CUDA capability information",
    )
    action.add_argument(
        "--smoke",
        choices=("cpu", "cuda"),
        help="run a tensor smoke test on exactly the selected device",
    )
    return parser


def main(argv: Sequence[str] | None = None) -> int:
    args = build_parser().parse_args(argv)
    try:
        if args.runtime_info:
            payload = runtime_info()
        elif args.smoke:
            payload = smoke(args.smoke)
        else:
            build_parser().print_help()
            return 0
    except RuntimeError as exc:
        print(json.dumps({"status": "error", "message": str(exc)}, ensure_ascii=False))
        return 2

    print(json.dumps(payload, ensure_ascii=False, sort_keys=True))
    return 0


if __name__ == "__main__":
    raise SystemExit(main())