"""Native vLLM draft_model speculative smoke + benchmark.
Target = Qwen3-VL bf16, draft = AWQ int4 self-twin (vision embeds shared from target).
--mode spec    : speculative engine, records token ids + timings + spec metrics
--mode vanilla : same engine settings without speculation (identity + speed reference)
Outputs one jsonl per run; identity check compares the two files offline.
"""
import argparse, io, json, time
from PIL import Image
import pyarrow.parquet as pq


def load_samples(parquet, pids_from, n, img_max_side):
    pids = []
    for line in open(pids_from):
        d = json.loads(line)
        if "pid" in d:
            pids.append(d["pid"])
        if len(pids) >= n:
            break
    tbl = pq.read_table(parquet, columns=["pid", "query", "decoded_image"])
    rowmap = {r["pid"]: r for r in tbl.to_pylist()}
    out = []
    for pid in pids:
        row = rowmap[pid]
        img = Image.open(io.BytesIO(row["decoded_image"]["bytes"])).convert("RGB")
        if max(img.size) > img_max_side:
            sc = img_max_side / max(img.size)
            img = img.resize((int(img.width * sc), int(img.height * sc)))
        out.append((pid, row["query"], img))
    return out


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--mode", choices=["spec", "vanilla"], required=True)
    ap.add_argument("--model", required=True)
    ap.add_argument("--draft", default="")
    ap.add_argument("--gamma", type=int, default=6)
    ap.add_argument("--parquet", required=True)
    ap.add_argument("--pids-from", required=True)
    ap.add_argument("--n", type=int, default=3)
    ap.add_argument("--max-tokens", type=int, default=256)
    ap.add_argument("--gpu-mem-util", type=float, default=0.35)
    ap.add_argument("--max-model-len", type=int, default=6144)
    ap.add_argument("--img-max-side", type=int, default=1024)
    ap.add_argument("--min-tokens", type=int, default=0)
    ap.add_argument("--out", required=True)
    ap.add_argument("--enforce-eager", action="store_true")
    ap.add_argument("--no-image", action="store_true")
    ap.add_argument("--tp", type=int, default=1)
    args = ap.parse_args()

    samples = load_samples(args.parquet, args.pids_from, args.n, args.img_max_side)
    from stage0_gonogo import build_inputs
    from transformers import AutoProcessor
    proc = AutoProcessor.from_pretrained(args.model)
    prompts = []
    for pid, query, img in samples:
        if args.no_image:
            img = None
        enc = build_inputs(proc, query, img)
        prompts.append((pid, enc["input_ids"][0].tolist(), img))
    print(f"[init] mode={args.mode} n={len(prompts)}", flush=True)

    from vllm import LLM, SamplingParams
    kw = {}
    if args.mode == "spec":
        kw["speculative_config"] = {
            "method": "draft_model",
            "model": args.draft,
            "num_speculative_tokens": args.gamma,
            "draft_tensor_parallel_size": args.tp,
        }
    llm = LLM(model=args.model, tensor_parallel_size=args.tp,
              disable_custom_all_reduce=True,
              gpu_memory_utilization=args.gpu_mem_util,
              max_model_len=args.max_model_len, enable_prefix_caching=False,
              enforce_eager=args.enforce_eager, disable_log_stats=False,
              max_num_batched_tokens=2048, max_num_seqs=4,
              limit_mm_per_prompt={"image": 1}, **kw)
    sp = SamplingParams(temperature=0, max_tokens=args.max_tokens,
                        min_tokens=args.min_tokens)

    def make_input(ids, img):
        if img is None:
            return {"prompt_token_ids": ids}
        return {"prompt_token_ids": ids, "multi_modal_data": {"image": img}}

    # warmup (not recorded)
    llm.generate([make_input(*prompts[0][1:])],
                 SamplingParams(temperature=0, max_tokens=8), use_tqdm=False)

    fout = open(args.out, "w")
    for pid, ids, img in prompts:
        t0 = time.time()
        out = llm.generate([make_input(ids, img)], sp, use_tqdm=False)[0]
        dt = time.time() - t0
        toks = list(out.outputs[0].token_ids)
        rec = dict(pid=pid, mode=args.mode, gamma=args.gamma if args.mode == "spec" else 0,
                   in_len=len(ids), gen_len=len(toks), t=round(dt, 3),
                   tok_s=round(len(toks) / dt, 2), gen_ids=toks)
        fout.write(json.dumps(rec) + "\n")
        fout.flush()
        print(f"[{pid}] gen={len(toks)} t={dt:.2f}s tok_s={len(toks)/dt:.1f}", flush=True)
    fout.close()

    try:
        for m in llm.get_metrics():
            if "spec" in m.name:
                val = getattr(m, "value", None)
                if val is None:
                    val = getattr(m, "values", None)
                print(f"[metric] {m.name} = {val}", flush=True)
    except Exception as e:
        print(f"[metric] unavailable: {e}", flush=True)
    print("[done]", flush=True)


if __name__ == "__main__":
    main()