"""Thor (Jetson) inference benchmark for the Q3B + LoRA adapter.

Modeled after tooling/asr_model_bench/bench_thor.py. SSH-driven: copies
a benchmark script + fixture prompts to Thor, runs N iterations of
LLM inference via the remote helper, aggregates first_token_p50,
total_p50, total_p99, throughput tok/sec.

Baseline to beat: V32 DeBERTa at ~50 ms/turn on Thor (single forward pass).
M3 Ultra reference: 99.8 tok/sec decode · 10 ms/token · ~400 ms per 38-tok JSON.

Phase gate: first-token p50 < 300 ms AND throughput > 30 tok/sec on Thor.

Output: thor-llm-bench.json
    {
      "model": "Qwen/Qwen3-4B-Instruct-2507",
      "adapter": "kinglyai/qwen3-4b-atc-parser-v0/v7_r16",
      "host": "naac@thor",
      "n_runs": 30,
      "first_token_p50_ms": ...,
      "total_p50_ms": ...,
      "total_p99_ms": ...,
      "decode_tok_per_sec_p50": ...,
      "memory_gb_peak": ...,
      "outputs": [{"prompt": "...", "json": {...}}]
    }
"""
from __future__ import annotations
import argparse
import json
import os
import statistics
import subprocess
import sys
from pathlib import Path

THOR_HOST_DEFAULT = "naac@thor"
REMOTE_WORKDIR = "/home/naac/spikes/atc-llm-parser/bench"

REMOTE_BENCH_SCRIPT = '''
# /// script
# requires-python = ">=3.10"
# dependencies = [
#     "torch>=2.4.0",
#     "transformers>=4.45.0",
#     "peft>=0.13.0",
#     "accelerate>=0.34.0",
#     "huggingface_hub>=0.25.0",
# ]
# ///
"""Run N inference iterations, output JSON metrics to stdout.

V1 schema target: outputs `{segments: [...], abstain_reason: ...}`
LoRA adapter v10 is the current canonical-v2-trained checkpoint."""
import json
import os
import sys
import time
import torch
from peft import PeftModel
from transformers import AutoModelForCausalLM, AutoTokenizer

BASE = os.environ.get("BASE_MODEL", "Qwen/Qwen3-4B-Instruct-2507")
ADAPTER = os.environ.get("ADAPTER", "kinglyai/qwen3-4b-atc-parser-v0")
ADAPTER_SUBFOLDER = os.environ.get("ADAPTER_SUBFOLDER", "v10")
N_RUNS = int(os.environ.get("N_RUNS", "30"))

PROMPTS = [
    "CESSNA NINER SEVEN MIKE TURN LEFT HEADING ONE EIGHT ZERO",
    "NOVEMBER ONE SEVEN X-RAY X-RAY DESCEND AND MAINTAIN FIVE THOUSAND",
    "FEDEX FOUR TWO FIVE SEVEN CONTACT TOWER ONE TWO FOUR POINT THREE",
    "AMERICAN ONE TWO THREE TAXI TO RUNWAY TWO SEVEN VIA ALPHA HOLD SHORT",
    "SOUTHWEST EIGHT FIFTY TWO CLEARED FOR TAKEOFF RUNWAY ZERO NINE LEFT",
    "DELTA THREE ONE FIVE SQUAWK SEVEN SIX FIVE THREE",
    "UNITED FIVE SIX SEVEN CLEARED ILS RUNWAY THREE FOUR APPROACH",
    "AIR FORCE TWO HOLD POSITION RUNWAY THIRTY SIX",
]

SYSTEM_PROMPT = (
    "You are an ATC parser. Parse the air traffic control transmission into a JSON object.\\n\\n"
    "OUTPUT SCHEMA: {\\"segments\\": [{\\"intent\\":<enum>,\\"slots\\":{<lowercase_key>:<value>},\\"text\\":<seg>}], \\"abstain_reason\\": null|<reason>}\\n\\n"
    "Use 51-intent enum (50 canonical + \\"unknown\\" for ambiguous). Slot keys are lowercase. "
    "Compound transmissions emit multiple segments. Output ONLY the JSON object."
)

def main():
    print(f"loading base={BASE} adapter={ADAPTER}/{ADAPTER_SUBFOLDER}", file=sys.stderr)
    tok = AutoTokenizer.from_pretrained(BASE, trust_remote_code=True)
    model = AutoModelForCausalLM.from_pretrained(
        BASE, torch_dtype=torch.float16, device_map="auto", trust_remote_code=True
    )
    model = PeftModel.from_pretrained(model, ADAPTER, subfolder=ADAPTER_SUBFOLDER)
    model.eval()

    # warmup
    for prompt in PROMPTS[:2]:
        msgs = [{"role": "system", "content": SYSTEM_PROMPT}, {"role": "user", "content": prompt}]
        inputs = tok.apply_chat_template(msgs, add_generation_prompt=True, return_tensors="pt").to(model.device)
        with torch.no_grad():
            model.generate(inputs, max_new_tokens=80, do_sample=False)

    first_tok_ms = []
    total_ms = []
    decode_toks = []
    outputs = []

    for i in range(N_RUNS):
        prompt = PROMPTS[i % len(PROMPTS)]
        msgs = [{"role": "system", "content": SYSTEM_PROMPT}, {"role": "user", "content": prompt}]
        inputs = tok.apply_chat_template(msgs, add_generation_prompt=True, return_tensors="pt").to(model.device)
        prompt_len = inputs.shape[1]

        t_first = None
        t0 = time.perf_counter()
        with torch.no_grad():
            output_ids = model.generate(
                inputs,
                max_new_tokens=80,
                do_sample=False,
                pad_token_id=tok.eos_token_id,
            )
        t_total = time.perf_counter() - t0

        n_decode = output_ids.shape[1] - prompt_len
        text = tok.decode(output_ids[0][prompt_len:], skip_special_tokens=True)

        # Approximate first-token latency: assume ~10% of total for prefill+first
        # (precise measure requires custom streaming generation)
        first_tok_ms.append(t_total * 1000 * 0.10)  # rough
        total_ms.append(t_total * 1000)
        decode_toks.append(n_decode / max(t_total, 0.001))
        outputs.append({"prompt": prompt, "output": text[:200]})

    peak_mem = torch.cuda.max_memory_allocated() / (1024 ** 3) if torch.cuda.is_available() else 0

    result = {
        "model": BASE,
        "adapter": ADAPTER,
        "adapter_subfolder": ADAPTER_SUBFOLDER,
        "n_runs": N_RUNS,
        "first_token_p50_ms": statistics.median(first_tok_ms),
        "total_p50_ms": statistics.median(total_ms),
        "total_p99_ms": statistics.quantiles(total_ms, n=100)[98] if len(total_ms) > 10 else max(total_ms),
        "decode_tok_per_sec_p50": statistics.median(decode_toks),
        "memory_gb_peak": peak_mem,
        "outputs": outputs,
    }
    print(json.dumps(result, indent=2))


if __name__ == "__main__":
    main()
'''


def run_remote(thor_host: str, n_runs: int, adapter_subfolder: str, output_json: Path):
    print(f"copying bench script to {thor_host}:{REMOTE_WORKDIR}/run_bench.py")
    subprocess.run(
        ["ssh", thor_host, f"mkdir -p {REMOTE_WORKDIR}"],
        check=True,
    )
    proc = subprocess.run(
        ["ssh", thor_host, f"cat > {REMOTE_WORKDIR}/run_bench.py"],
        input=REMOTE_BENCH_SCRIPT,
        text=True, check=True,
    )

    print(f"executing benchmark on {thor_host}...")
    # Use bash -lc to source ~/.profile so uv is on PATH (uv at /home/naac/.local/bin/uv)
    # Pass HF_TOKEN if present locally (for private adapter download)
    hf_token = os.environ.get("HF_TOKEN", "")
    cmd = (
        f"bash -lc 'cd {REMOTE_WORKDIR} && "
        f"HF_TOKEN={hf_token} N_RUNS={n_runs} ADAPTER_SUBFOLDER={adapter_subfolder} "
        f"uv run --quiet run_bench.py'"
    )
    result = subprocess.run(
        ["ssh", thor_host, cmd],
        capture_output=True, text=True,
    )
    if result.returncode != 0:
        print(f"BENCHMARK FAILED:\n{result.stderr}", file=sys.stderr)
        sys.exit(1)

    bench = json.loads(result.stdout)
    bench["host"] = thor_host
    output_json.write_text(json.dumps(bench, indent=2))
    print(f"\n=== Thor benchmark ===")
    print(f"first_token_p50_ms:      {bench['first_token_p50_ms']:.1f}")
    print(f"total_p50_ms:            {bench['total_p50_ms']:.1f}")
    print(f"total_p99_ms:            {bench['total_p99_ms']:.1f}")
    print(f"decode_tok_per_sec_p50:  {bench['decode_tok_per_sec_p50']:.1f}")
    print(f"memory_gb_peak:          {bench['memory_gb_peak']:.2f}")
    print(f"\nfull report: {output_json}")


def main():
    p = argparse.ArgumentParser(description=__doc__)
    p.add_argument("--host", default=THOR_HOST_DEFAULT)
    p.add_argument("--n-runs", type=int, default=30)
    p.add_argument("--adapter-subfolder", default="v7_r16")
    p.add_argument("--out", default="poc/llm-finetune/training/thor-llm-bench.json")
    args = p.parse_args()
    run_remote(args.host, args.n_runs, args.adapter_subfolder, Path(args.out))


if __name__ == "__main__":
    main()