"""Comprehensive GPU validation for the ASD custom generation method.

Covers:
1. Native strict baseline vs ASD(B=0)        -> must be token-identical (gold check)
2. Native strict vs custom (no knobs)        -> must be token-identical (lossless default)
3. ASD(B=2.0) relaxed                        -> runs, deterministic, >= strict acceptance
4. Determinism                               -> two identical relaxed runs, token-identical
5. Error handling                            -> do_sample / ensemble_weight / no draft source
6. Throughput + acceptance comparison        -> strict vs relaxed
"""

import os
import sys
import time

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

REPO = os.environ.get("ASD_REPO", os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
TARGET = os.environ.get("ASD_TARGET", "Qwen/Qwen2.5-1.5B-Instruct")
DRAFT = os.environ.get("ASD_DRAFT", "Qwen/Qwen2.5-0.5B-Instruct")
MAX_NEW = 128
PROMPTS = [
    "Give me a short introduction to large language models.",
    "What is the capital of France? Answer in one sentence.",
    "Write a Python function that computes the nth Fibonacci number.",
    "Explain the difference between TCP and UDP briefly.",
    "List three prime numbers greater than 10 and justify one of them.",
]

failures = []


def check(name, cond, detail=""):
    status = "PASS" if cond else "FAIL"
    print(f"[{status}] {name}" + (f" -- {detail}" if detail else ""))
    if not cond:
        failures.append(name)


def count_fwd(tok, out, prompt_len):
    return out.shape[1] - prompt_len


def main():
    tokenizer = AutoTokenizer.from_pretrained(TARGET)
    model = AutoModelForCausalLM.from_pretrained(TARGET, dtype="auto", device_map="auto")
    assistant = AutoModelForCausalLM.from_pretrained(DRAFT, dtype="auto", device_map="auto")
    model.eval()
    assistant.eval()

    def run(prompt, **extra):
        inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
        t0 = time.perf_counter()
        out = model.generate(
            **inputs, max_new_tokens=MAX_NEW, do_sample=False, **extra
        )
        dt = time.perf_counter() - t0
        n = out.shape[1] - inputs["input_ids"].shape[1]
        return out, n / dt

    # ---------- 1&2. gold checks ----------
    for prompt in PROMPTS:
        strict, _ = run(prompt, assistant_model=assistant)
        b0, _ = run(
            prompt,
            assistant_model=assistant,
            custom_generate=REPO,
            trust_remote_code=True,
            assistant_asd_budget=0.0,
        )
        noknobs, _ = run(
            prompt,
            assistant_model=assistant,
            custom_generate=REPO,
            trust_remote_code=True,
        )
        check(f"gold B=0 identical ({prompt[:25]!r})", torch.equal(strict, b0))
        check(f"gold no-knobs identical ({prompt[:25]!r})", torch.equal(strict, noknobs))

    # ---------- 3&4. relaxed run + determinism ----------
    for prompt in PROMPTS:
        relaxed1, tps1 = run(
            prompt,
            assistant_model=assistant,
            custom_generate=REPO,
            trust_remote_code=True,
            assistant_asd_budget=2.0,
            assistant_asd_local_ratio=0.25,
            assistant_asd_max_mismatches=2,
        )
        relaxed2, _ = run(
            prompt,
            assistant_model=assistant,
            custom_generate=REPO,
            trust_remote_code=True,
            assistant_asd_budget=2.0,
            assistant_asd_local_ratio=0.25,
            assistant_asd_max_mismatches=2,
        )
        check(f"determinism ({prompt[:25]!r})", torch.equal(relaxed1, relaxed2))
        check(f"relaxed non-empty ({prompt[:25]!r})", relaxed1.shape[1] > 0)

    # ---------- 5. error handling ----------
    inputs = tokenizer(PROMPTS[0], return_tensors="pt").to(model.device)

    def expect_raise(name, **kw):
        try:
            model.generate(**inputs, max_new_tokens=8, **kw)
        except ValueError:
            check(f"raises {name}", True)
        except Exception as e:  # noqa: BLE001
            check(f"raises {name}", False, f"wrong exception: {type(e).__name__}: {e}")
        else:
            check(f"raises {name}", False, "no exception")

    expect_raise(
        "do_sample=True",
        assistant_model=assistant,
        custom_generate=REPO,
        trust_remote_code=True,
        do_sample=True,
        assistant_asd_budget=2.0,
    )
    expect_raise(
        "ensemble_weight conflict",
        assistant_model=assistant,
        custom_generate=REPO,
        trust_remote_code=True,
        assistant_ensemble_weight=0.5,
        assistant_asd_budget=2.0,
    )
    expect_raise(
        "knobs without draft source",
        custom_generate=REPO,
        trust_remote_code=True,
        assistant_asd_budget=2.0,
    )
    expect_raise(
        "negative budget",
        assistant_model=assistant,
        custom_generate=REPO,
        trust_remote_code=True,
        assistant_asd_budget=-1.0,
    )

    # ---------- 6. throughput + acceptance ----------
    print("\n=== throughput comparison ===")
    tot_strict_tps = tot_relaxed_tps = 0.0
    for prompt in PROMPTS:
        strict, s_tps = run(prompt, assistant_model=assistant)
        relaxed, r_tps = run(
            prompt,
            assistant_model=assistant,
            custom_generate=REPO,
            trust_remote_code=True,
            assistant_asd_budget=2.0,
            assistant_asd_local_ratio=0.25,
            assistant_asd_max_mismatches=2,
        )
        tot_strict_tps += s_tps
        tot_relaxed_tps += r_tps
        overlap = sum(
            1 for a, b in zip(strict[0].tolist(), relaxed[0].tolist()) if a == b
        ) / max(len(strict[0]), len(relaxed[0]))
        print(
            f"  strict {s_tps:6.2f} tok/s | relaxed {r_tps:6.2f} tok/s "
            f"({(r_tps / s_tps - 1) * 100:+5.1f}%) | prefix overlap {overlap:.2%}"
        )
    print(
        f"  MEAN strict {tot_strict_tps / len(PROMPTS):.2f} tok/s | "
        f"relaxed {tot_relaxed_tps / len(PROMPTS):.2f} tok/s"
    )

    print("\n" + ("ALL PASS" if not failures else f"FAILURES: {failures}"))
    sys.exit(1 if failures else 0)


if __name__ == "__main__":
    main()