"""Call every method through the /trace API and render each trace with the UI's own code.

Usage:
  python scripts/smoke_test.py --url http://127.0.0.1:7860 --families qwen25,smollm2
  python scripts/smoke_test.py --space USER/SPACE --families qwen3,llama32,gemma3,smollm2 --tokens 16
"""

from __future__ import annotations

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

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
os.environ["DD_SKIP_INIT"] = "1"  # import app.py for its renderers without loading models

from viz import plots  # noqa: E402

CASES = [
    ("contrastive", {"method": "cd", "prompt": "Answer in one word: what is the capital of Australia?", "alpha": 1.0,
                     "beta": 0.1}),
    ("contrastive", {"method": "cad", "prompt": "Based on the passage, in which city does the Eiffel Tower stand?",
                     "context": "In 2031 the Eiffel Tower was rebuilt, piece by piece, in Rome.", "alpha": 1.0}),
    ("contrastive", {"method": "rose", "prompt": "Explain in two sentences why the sky is blue.", "alpha": 0.5}),
    ("contrastive", {"method": "dola", "prompt": "What is the capital of Washington State? Answer in one word.",
                     "rep": 1.2}),
    ("guided", {"method": "classifier", "attribute": "sentiment", "target": "positive", "lam": 3.0, "k": 8,
                "lookahead": 1, "prompt": "Write one sentence describing a rainy Monday morning."}),
    ("guided", {"method": "classifier", "attribute": "topic", "topic": "astronomy", "lam": 20.0, "k": 8,
                "prompt": "Write two sentences about a trip to the beach."}),
    ("guided", {"method": "regex", "pattern": r"\d{4}-\d{2}-\d{2}",
                "prompt": "On what date did the Berlin Wall fall? Reply with the date only, in YYYY-MM-DD format."}),
    ("parallel", {"method": "speculative", "gamma": 4, "mode": "greedy",
                  "prompt": "List the numbers from 1 to 40, separated by commas."}),
    ("parallel", {"method": "speculative", "gamma": 4, "mode": "sampling", "temperature": 0.8,
                  "prompt": "Write a haiku about quantum entanglement."}),
    ("parallel", {"method": "pld", "num_pred": 8,
                  "prompt": "Fix only the typos and return the text otherwise unchanged: 'The quick brown fox jumpd "
                            "over the lazzy dog and ran into the forrest.'"}),
    ("parallel", {"method": "jacobi", "block": 6, "prompt": "Count from 1 to 30, separated by spaces."}),
]


def render_all(trace: dict) -> None:
    """Exercise the same render/step functions the UI calls for this trace."""
    import app

    app.time.sleep = lambda _s: None  # replays animate with sleeps; skip them here
    if trace["paradigm"] == "contrastive":
        app.render_contrastive(trace)
        n = len(trace["runs"]["method"]["steps"])
        for i in {0, n // 2, max(n - 1, 0)}:
            app.step_contrastive(trace, i)
        list(app.replay_contrastive(trace))
    elif trace["paradigm"] == "guided":
        app.render_guided(trace)
        n = len(trace["runs"]["method"]["steps"])
        for i in {0, n // 2, max(n - 1, 0)}:
            app.step_guided(trace, i)
        list(app.replay_guided(trace))
    else:
        app.render_parallel(trace)
        n = len(trace["runs"]["method"]["iters"])
        for k in {0, n // 2, max(n - 1, 0)}:
            app.step_parallel(trace, k)
        list(app.replay_parallel(trace))
        plots.speedup_curves(trace["summary"] if trace["method"] == "speculative" else None, 0.7, 0.5)


def check(trace: dict) -> list[str]:
    problems = []
    if trace.get("schema") != "dd.trace/v1":
        problems.append("bad schema")
    meth = trace["runs"]["method"]
    if not meth["ids"]:
        problems.append("method produced no tokens")
    s = trace["summary"]
    if trace["paradigm"] == "parallel" and trace["params"]["mode"] == "greedy" and s.get("identical") is False:
        problems.append(f"greedy output diverged from autoregressive at token {s.get('first_divergence')} (bf16?)")
    if trace["method"] == "regex" and not s["valid"]["method"] and meth["finish"] != "length":
        problems.append(f"regex output does not match: {meth['text']!r}")
    json.dumps(trace, allow_nan=False)
    return problems


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--url")
    ap.add_argument("--space")
    ap.add_argument("--families", default="qwen25,smollm2")
    ap.add_argument("--tokens", type=int, default=24)
    ap.add_argument("--only", default="", help="comma-separated methods to run")
    args = ap.parse_args()

    from gradio_client import Client

    from huggingface_hub import get_token

    token = os.environ.get("HF_TOKEN") or get_token()  # logged-in calls get the account's GPU quota
    client = Client(args.space, token=token) if args.space else Client(args.url or "http://127.0.0.1:7860")
    only = {m for m in args.only.split(",") if m}
    failures = 0
    for family in args.families.split(","):
        for paradigm, params in CASES:
            if only and params["method"] not in only:
                continue
            P = {**params, "max_new_tokens": args.tokens}
            t0 = time.time()
            try:
                trace = client.predict(paradigm, family, json.dumps(P), api_name="/trace")
                render_all(trace)
                problems = check(trace)
            except Exception as exc:  # report and keep going
                problems = [f"{type(exc).__name__}: {str(exc)[:300]}"]
                trace = None
            dt = time.time() - t0
            tag = "OK  " if not problems else "FAIL"
            failures += bool(problems)
            extra = ""
            if trace:
                s = trace["summary"]
                env = trace.get("env", {})
                extra = (f"gpu {env.get('gpu_seconds')}s/{env.get('requested_seconds')}s  "
                         + " ".join(f"{k}={s[k]}" for k in ("tau", "speedup_wall", "accept_rate", "changed", "identical")
                                    if k in s))
                extra += f"  | {trace['runs']['method']['text'][:70]!r}"
            print(f"{tag} {family:8} {paradigm:11} {params['method']:11} {dt: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())