"""Call the Space's /trace API and render every result with the UI's own code.

Usage:
  python scripts/smoke_test.py --url http://127.0.0.1:7860
  python scripts/smoke_test.py --space wang2226/steering-showcase [--models qwen25-05b,smollm2-135m] [--custom]

Every call that is not a cache hit spends GPU time on a Space; logged-in calls use the account's quota.
"""

from __future__ import annotations

import argparse
import os
import sys
import time
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
os.environ["SD_SKIP_INIT"] = "1"  # import the renderers without loading models

from content import examples as EX  # noqa: E402
from steering.core import ALPHAS  # noqa: E402
from steering.scenarios import SCENARIOS  # noqa: E402
from viz import render as R  # noqa: E402


def check(trace: dict, model: str, scenario: str) -> list[str]:
    problems = []
    if trace.get("schema") != "steer.trace/v1":
        problems.append("bad schema")
    if trace.get("source") != "live":
        problems.append("not a live trace")
    if [s["alpha"] for s in trace["states"]] != ALPHAS:
        problems.append("states off the alpha grid")
    if trace.get("direction") != scenario:
        problems.append(f"direction {trace.get('direction')!r}")
    for s in trace["states"]:
        if not s["continuation"].startswith(" ".join(trace["candidates"][s["selected_index"]].split())):
            problems.append(f"alpha {s['alpha']}: response does not start with the highlighted token")
        if abs(sum(s["percents"]) + s["other"] - 100) > 0.01:
            problems.append(f"alpha {s['alpha']}: probabilities do not sum to 100")
    if not trace.get("depth") or not trace.get("negative_examples"):
        problems.append("depth or contrast continuations missing")
    for alpha in ALPHAS:  # the same renderers the page uses
        R.formula_md(trace, alpha), R.cards_html(trace, alpha), R.response_html(trace, alpha)
        R.chart_html(trace, alpha), R.sweep_html(trace, alpha), R.schematic_html(trace, alpha)
    R.takeaway_html(trace), R.contrast_html(trace), R.explain_html(trace)
    return problems


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--url")
    ap.add_argument("--space")
    ap.add_argument("--models", default="qwen25-05b")
    ap.add_argument("--custom", action="store_true", help="also run one visitor-style prompt per model")
    ap.add_argument("--fresh", action="store_true", help="bypass the Space's cache")
    args = ap.parse_args()

    from gradio_client import Client
    from huggingface_hub import get_token

    token = os.environ.get("HF_TOKEN") or get_token()
    client = Client(args.space, token=token) if args.space else Client(args.url or "http://127.0.0.1:7860")
    cases = [(s.id, "", "") for s in SCENARIOS]
    if args.custom:
        e = EX.EXAMPLES[-1]
        cases.append((e["scenario"], e["prompt"], e["prefix"]))
    failures = 0
    for model in args.models.split(","):
        for scenario, prompt, prefix in cases:
            t0 = time.time()
            try:
                trace = client.predict(model, scenario, prompt, prefix, args.fresh, api_name="/trace")
                problems = check(trace, model, scenario)
            except Exception as exc:
                trace, problems = None, [f"{type(exc).__name__}: {str(exc)[:300]}"]
            failures += bool(problems)
            extra = ""
            if trace:
                env = trace.get("env", {})
                winners = " → ".join(dict.fromkeys(R.label(trace["candidates"][s["selected_index"]]) for s in trace["states"]))
                extra = (f"{'cached' if trace.get('cached') else 'fresh'} {env.get('device')} {env.get('gpu_seconds')}s"
                         f" | {winners}")
            print(f"{'OK  ' if not problems else 'FAIL'} {model:13} {scenario:18} {'custom' if prompt else '      '} "
                  f"{time.time() - t0:5.1f}s  {extra}")
            for p in problems:
                print(f"       ! {p}")
    print(f"\n{failures} failure(s)")
    return 1 if failures else 0


if __name__ == "__main__":
    sys.exit(main())