#!/usr/bin/env python3
"""
peer_run.py -- generate completions for peer-model NIAH prompts.

Backends, tried in this order unless --backend forces one:
  vllm   fast, batched. Cannot load Spark-X2.5 (model_type spark2_5 is
         custom_code and has no vLLM backend), so it is skipped automatically.
  hf     transformers with trust_remote_code. Works for everything, slower.

Prompts are fed as TOKEN IDS, never text: peer_gen.py guarantees an exact token
count and a decode->re-encode round trip would silently break it.

Greedy decoding, matching our own harness (CP_GEN_GREEDY=1). Generation stops at
--max-new tokens; no early stop, because the peer models have no shared
end-of-answer convention and truncation is harmless -- scoring is substring
match, so trailing chatter costs nothing.

  usage:
    peer_run.py --model Qwen/Qwen3-4B-Base --prompts out/niah_single_prompts.jsonl \
                --out gen/qwen3-4b_niah_single.jsonl --max-new 24
"""

import argparse
import json
import os
import sys
import time


def load_prompts(p):
    rows = []
    with open(p) as f:
        for line in f:
            line = line.strip()
            if line:
                rows.append(json.loads(line))
    return rows


def run_vllm(model, rows, max_new, gpu_frac, max_len):
    from vllm import LLM, SamplingParams
    llm = LLM(model=model, dtype="bfloat16", trust_remote_code=True,
              gpu_memory_utilization=gpu_frac, max_model_len=max_len,
              enforce_eager=False)
    sp = SamplingParams(temperature=0.0, max_tokens=max_new)
    outs = llm.generate(prompt_token_ids=[r["ids"] for r in rows],
                        sampling_params=sp)
    return [o.outputs[0].text for o in outs]


def run_hf(model, rows, max_new, batch):
    import torch
    from transformers import AutoModelForCausalLM, AutoTokenizer
    tok = AutoTokenizer.from_pretrained(model, trust_remote_code=True)
    m = AutoModelForCausalLM.from_pretrained(
        model, dtype=torch.bfloat16, device_map="cuda", trust_remote_code=True)
    m.eval()
    texts = []
    # batch=1 by default: these are 16k-token prompts and padding a ragged batch
    # would change each sequence's positions, which is exactly what NIAH measures.
    for s in range(0, len(rows), batch):
        chunk = rows[s:s + batch]
        ids = torch.tensor([r["ids"] for r in chunk], device="cuda")
        with torch.no_grad():
            out = m.generate(input_ids=ids, max_new_tokens=max_new,
                             do_sample=False, num_beams=1,
                             pad_token_id=getattr(tok, "pad_token_id", None)
                             or getattr(tok, "eos_token_id", 0))
        for j in range(len(chunk)):
            texts.append(tok.decode(out[j][ids.shape[1]:], skip_special_tokens=True))
        done = min(s + batch, len(rows))
        print(f"    {done}/{len(rows)}", flush=True)
    return texts


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--model", required=True)
    ap.add_argument("--prompts", required=True)
    ap.add_argument("--out", required=True)
    ap.add_argument("--max-new", type=int, default=24)
    ap.add_argument("--backend", default="auto", choices=["auto", "vllm", "hf"])
    ap.add_argument("--gpu-frac", type=float, default=0.90)
    ap.add_argument("--batch", type=int, default=1, help="hf backend only")
    args = ap.parse_args()

    rows = load_prompts(args.prompts)
    if not rows:
        sys.exit(f"no prompts in {args.prompts}")
    plen = len(rows[0]["ids"])
    max_len = plen + args.max_new + 8
    # Never ask vLLM for more than the checkpoint's trained window: Qwen3-Base
    # ships max_position_embeddings=32768, so a 32,744-token prompt + 24 new + 8
    # slack would exceed it and vLLM refuses to start. Clamp, and say so.
    try:
        from transformers import AutoConfig
        mp = getattr(AutoConfig.from_pretrained(args.model, trust_remote_code=True),
                     "max_position_embeddings", None)
        if mp and max_len > mp:
            print(f"  clamping max_model_len {max_len} -> {mp} (model's native window)")
            max_len = mp
            if plen + args.max_new > mp:
                sys.exit(f"  FATAL: prompt {plen} + max_new {args.max_new} exceeds "
                         f"the model's window {mp}; lower --prompt-len")
    except Exception as e:
        print(f"  (could not read max_position_embeddings: {type(e).__name__})")
    print(f"  {args.model}  {len(rows)} prompts x {plen} tok  max_new={args.max_new}")

    backend = args.backend
    if backend == "auto":
        backend = "hf"
        try:
            import vllm  # noqa: F401
            # spark2_5 is custom_code with no vLLM backend; do not waste a load
            backend = "hf" if "spark" in args.model.lower() else "vllm"
        except ImportError:
            pass
    print(f"  backend: {backend}")

    t0 = time.time()
    if backend == "vllm":
        try:
            texts = run_vllm(args.model, rows, args.max_new, args.gpu_frac, max_len)
        except Exception as e:
            print(f"  vLLM failed ({type(e).__name__}: {str(e)[:140]}) -- falling back to hf")
            texts = run_hf(args.model, rows, args.max_new, args.batch)
    else:
        texts = run_hf(args.model, rows, args.max_new, args.batch)
    dt = time.time() - t0

    os.makedirs(os.path.dirname(args.out) or ".", exist_ok=True)
    with open(args.out, "w") as f:
        for r, t in zip(rows, texts):
            f.write(json.dumps({"i": r["i"], "generated": t}) + "\n")
    print(f"  wrote {args.out}  ({len(texts)} rows, {dt:.0f}s, "
          f"{dt/max(1,len(texts)):.2f}s/prompt)")


if __name__ == "__main__":
    main()