from __future__ import annotations

import argparse
import json
import shutil
import subprocess
import sys
import urllib.error
import urllib.request
from dataclasses import asdict, dataclass
from pathlib import Path
from typing import Any, Literal


ROOT = Path(__file__).resolve().parents[1]
if str(ROOT) not in sys.path:
    sys.path.insert(0, str(ROOT))

DEFAULT_BASE_URL = "http://127.0.0.1:8080/v1"
DEFAULT_MODEL = "signal-garden-qwen35-code-mutation"
REQUIRED_LLAMA_HELP_TERMS = (
    "-hf",
    "--alias",
    "--ctx-size",
    "--host",
    "--port",
    "--reasoning",
    "--reasoning-budget",
)

Status = Literal["ok", "warn", "fail"]


@dataclass(frozen=True)
class CheckResult:
    name: str
    status: Status
    detail: str


def check_python_app() -> CheckResult:
    try:
        import gradio as gr
        import numpy as np

        import app

        app.build_demo(app.ArcadeConfig())
    except Exception as exc:
        return CheckResult("python-app", "fail", f"Gradio app import/build failed: {exc}")

    return CheckResult(
        "python-app",
        "ok",
        f"Gradio {getattr(gr, '__version__', 'unknown')} and NumPy {np.__version__} import; app builds.",
    )


def check_uv() -> CheckResult:
    uv_path = shutil.which("uv")
    if uv_path is None:
        return CheckResult("uv", "fail", "`uv` is not on PATH; repository commands use `uv run ...`.")
    return CheckResult("uv", "ok", f"`uv` found at {uv_path}.")


def check_llama_server() -> CheckResult:
    llama_path = shutil.which("llama-server")
    if llama_path is None:
        return CheckResult(
            "llama-server",
            "fail",
            "`llama-server` is not on PATH; install llama.cpp or use the Docker Space build.",
        )

    try:
        completed = subprocess.run(
            [llama_path, "--help"],
            check=False,
            capture_output=True,
            text=True,
            timeout=15,
        )
    except subprocess.TimeoutExpired:
        return CheckResult("llama-server", "fail", f"`{llama_path} --help` timed out.")
    except OSError as exc:
        return CheckResult("llama-server", "fail", f"Could not run `{llama_path} --help`: {exc}")

    help_text = f"{completed.stdout}\n{completed.stderr}"
    if completed.returncode != 0:
        return CheckResult("llama-server", "fail", f"`{llama_path} --help` exited {completed.returncode}.")

    missing = [term for term in REQUIRED_LLAMA_HELP_TERMS if term not in help_text]
    if missing:
        return CheckResult(
            "llama-server",
            "fail",
            f"`{llama_path}` is missing required llama.cpp server flag(s): {', '.join(missing)}.",
        )

    return CheckResult(
        "llama-server",
        "ok",
        f"`llama-server` found at {llama_path} with HF, alias, context, host/port, and reasoning flags.",
    )


def check_llm_endpoint(base_url: str, model: str, timeout_seconds: float, required: bool) -> CheckResult:
    endpoint = f"{base_url.rstrip('/')}/models"
    failure_status: Status = "fail" if required else "warn"
    try:
        with urllib.request.urlopen(endpoint, timeout=max(0.1, float(timeout_seconds))) as response:
            body = response.read().decode("utf-8")
    except urllib.error.URLError as exc:
        return CheckResult(
            "llm-endpoint",
            failure_status,
            f"{endpoint} is not reachable: {exc.reason}. Start llama-server before playing LLM rounds.",
        )
    except TimeoutError:
        return CheckResult(
            "llm-endpoint",
            failure_status,
            f"{endpoint} timed out after {timeout_seconds:g}s.",
        )

    try:
        payload = json.loads(body)
    except json.JSONDecodeError as exc:
        return CheckResult("llm-endpoint", failure_status, f"{endpoint} returned malformed JSON: {exc.msg}.")

    ids = _model_ids(payload)
    if ids and model not in ids:
        return CheckResult(
            "llm-endpoint",
            failure_status,
            f"{endpoint} is live, but model alias `{model}` was not listed. Available: {', '.join(ids)}.",
        )
    if not ids:
        return CheckResult("llm-endpoint", "warn", f"{endpoint} is live, but no model ids were listed.")

    return CheckResult("llm-endpoint", "ok", f"{endpoint} is live and lists `{model}`.")


def run_checks(
    base_url: str = DEFAULT_BASE_URL,
    model: str = DEFAULT_MODEL,
    timeout_seconds: float = 5.0,
    require_live_llm: bool = False,
) -> list[CheckResult]:
    return [
        check_python_app(),
        check_uv(),
        check_llama_server(),
        check_llm_endpoint(base_url, model, timeout_seconds, require_live_llm),
    ]


def _model_ids(payload: Any) -> tuple[str, ...]:
    if not isinstance(payload, dict):
        return ()
    data = payload.get("data")
    if not isinstance(data, list):
        return ()
    ids: list[str] = []
    for item in data:
        if isinstance(item, dict) and isinstance(item.get("id"), str):
            ids.append(item["id"])
    return tuple(ids)


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(description="Check Signal Garden Gradio and LLM runtime requirements.")
    parser.add_argument("--base-url", default=DEFAULT_BASE_URL)
    parser.add_argument("--model", default=DEFAULT_MODEL)
    parser.add_argument("--timeout-seconds", type=float, default=5.0)
    parser.add_argument(
        "--require-live-llm",
        action="store_true",
        help="Fail when the OpenAI-compatible /models endpoint is not reachable or has the wrong model alias.",
    )
    parser.add_argument("--json", action="store_true", help="Print machine-readable check results.")
    args = parser.parse_args(argv)

    results = run_checks(
        base_url=args.base_url,
        model=args.model,
        timeout_seconds=max(0.1, float(args.timeout_seconds)),
        require_live_llm=bool(args.require_live_llm),
    )
    if args.json:
        print(json.dumps([asdict(result) for result in results], indent=2, sort_keys=True))
    else:
        for result in results:
            print(f"{result.status.upper():4} {result.name}: {result.detail}")

    return 1 if any(result.status == "fail" for result in results) else 0


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