"""Run the tiny V4.1 checkpoint through the real DeepInfer engine, teacher-forced on the reference golden.

Engine path: deepinfer.entrypoints.llm.LLM (scheduler, KV cache, model runner, sampler), EP=DP=CP=1, Neuron basic
mode defaults (eager, DEEPINFER_USE_NEXTN_DECODE=0). Teacher forcing uses the engine's own mechanism,
SamplingParams(golden_token_ids=...), so decode step i is fed the golden token i-1 exactly like the reference
(reference_golden.py --force-tokens-dir). Logits are captured by wrapping the sampler's `get_logits` (the fp32
[num_reqs, vocab] lm-head output before any sampling transform); per-layer block outputs by forward hooks on the
decoder layers (best effort, see --hidden).

Output: <out>/<prompt>.pt in the reference_golden.py format (prompt_ids, tokens = argmax per step, fed_tokens,
logits [steps+1, vocab] fp32, optional hidden), ready for compare_golden.py.

    cd ~/di/DeepInfer && ~/di/venv/bin/python examples/neuron/run_deepinfer_tiny.py \
        --hf-ckpt ~/di/tiny_ckpt --golden-dir ~/di/golden --out ~/di/deepinfer_out [--device cpu|neuron] \
        [--hidden] [--compare] [--threads 32]

Placeholder engine steps (is_placeholder=True, no real sequence) also run the model and the sampler; they are
filtered out of the recorded logits/hidden. `install_workarounds` patches, at runtime and only in this process,
the bugs found in other modules while bringing this up (BUG-1/2/4, NOTE-3/5 in the code); the list of the ones
that actually triggered is printed at the end. Set RUN_TINY_DEBUG=1 to print every sampler call.
"""

from __future__ import annotations

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

REPO_ROOT = Path(__file__).resolve().parents[2]


def parse_args():
    ap = argparse.ArgumentParser()
    ap.add_argument("--hf-ckpt", default="~/di/tiny_ckpt")
    ap.add_argument("--golden-dir", default="~/di/golden")
    ap.add_argument("--out", default="~/di/deepinfer_out")
    ap.add_argument("--prompts", default="short,medium,long")
    ap.add_argument("--device", choices=["cpu", "neuron"], default="cpu", help="sets DEEPINFER_NEURON_DEVICE")
    ap.add_argument("--max-model-len", type=int, default=2048)
    ap.add_argument("--max-num-batched-tokens", type=int, default=1024)
    ap.add_argument("--threads", type=int, default=16)
    ap.add_argument("--num-kv-blocks", type=int, default=32, help="fixed KV block pool (128 tokens each); 0 = engine profiling")
    ap.add_argument("--hidden", action="store_true", help="also record per-layer block outputs (hc residual stream)")
    ap.add_argument("--compare", action="store_true", help="run compare_golden on every prompt afterwards")
    return ap.parse_args()


def setup_env(device: str) -> None:
    # Must happen before deepinfer is imported (see examples/_config.py).
    os.environ["DEEPINFER_ACCELERATOR"] = "neuron"
    os.environ["DEEPINFER_NEURON_DEVICE"] = device
    os.environ["DEEPINFER_USE_NEXTN_DECODE"] = "0"
    os.environ.setdefault("DEEPINFER_ENFORCE_EAGER", "1")
    os.environ.setdefault("DEEPINFER_ZERO_OVERHEAD_SCHEDULING", "0")
    os.environ.setdefault("CUDA_VISIBLE_DEVICES", "0")  # one engine rank; mapped to the target device
    os.environ.setdefault("RANK", "0")
    os.environ.setdefault("WORLD_SIZE", "1")
    os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")
    os.environ.setdefault("DEEPINFER_USE_MULTIMODAL_POOL", "0")  # text only; avoids an 8 GB pinned host pool
    if str(REPO_ROOT) not in sys.path:
        sys.path.insert(0, str(REPO_ROOT))
    here = str(Path(__file__).resolve().parent)
    if here not in sys.path:
        sys.path.insert(0, here)


WORKAROUNDS: list[str] = []


def install_workarounds(num_kv_blocks: int) -> None:
    """Work around bugs in other modules without editing them; each one is reported (see the module docstring)."""
    import torch
    import torch.distributed as dist

    # BUG-1: PlatformNeuron(cpu) uses gloo and fp8_sf_type=torch.int16; gloo collectives reject int16
    # ("RuntimeError: Invalid scalar type") in v4_modules/engram.py NgramEmbedding.gather_sf. Byte-reinterpret
    # dtypes gloo lacks as uint8 for the gather-type collectives (lossless; concatenation is along dim 0).
    orig_agit = dist.all_gather_into_tensor
    unsupported = {torch.int16, torch.uint16, torch.float8_e4m3fn, torch.float8_e8m0fnu}

    def all_gather_into_tensor(output_tensor, input_tensor, group=None, async_op=False):
        if input_tensor.dtype in unsupported and dist.get_backend(group) == "gloo":
            if "BUG-1" not in WORKAROUNDS:
                WORKAROUNDS.append("BUG-1")
            return orig_agit(output_tensor.view(torch.uint8), input_tensor.contiguous().view(torch.uint8), group=group, async_op=async_op)
        return orig_agit(output_tensor, input_tensor, group=group, async_op=async_op)

    dist.all_gather_into_tensor = all_gather_into_tensor

    # BUG-2: model_executor/model_loader.py get_model asserts `p.is_cuda` for every parameter; on Neuron the
    # parameters live on the transfer_to_neuron target (neuron:0 or cpu). Same function, device-agnostic check.
    import deepinfer.model_executor.model_loader as ml
    import deepinfer.worker.model_runner as mr
    from deepinfer import PLATFORM

    def get_model(model_config, skip_load=False):
        model_class = ml._get_model_architecture(model_config)
        with ml._set_default_torch_dtype(model_config.dtype), torch.device("cuda"):
            model = model_class.load_model(model_config.model, model_config.model_cfg, skip_load)
        target = PLATFORM.device.type
        for n, p in model.named_parameters():
            if p.device.type != target:
                raise AssertionError(f"{n} on {p.device}, expected {target}")
            if hasattr(p, "sf") and p.sf.device.type != target:
                raise AssertionError(f"{n}.sf on {p.sf.device}, expected {target}")
        if not any("BUG-2" == w for w in WORKAROUNDS):
            WORKAROUNDS.append("BUG-2")
        return model.eval().requires_grad_(False)

    ml.get_model = get_model
    mr.get_model = get_model

    # NOTE-5: WorkerBase pins the worker to phys_cpus // get_actual_num_gpus() CPUs; on Trn2 that is 96 // 64 = 1
    # CPU, so CPU-mode runs are single-threaded (~10x slower). Keep the process affinity instead.
    import deepinfer.worker.base as wb

    def bind_cpu(local_rank):
        WORKAROUNDS.append("NOTE-5: skipped bind_cpu (1 CPU per rank on Trn2)")

    wb.bind_cpu = bind_cpu

    # BUG-4: vendor deep_gemm fp8_fp4_paged_mqa_logits gathers keys for every column < max_context_len and only
    # masks afterwards, so block-table slots past a request's last page (unused entries, not valid page ids) are
    # dereferenced -> IndexError. Sanitize out-of-range table entries (their columns are masked to -inf anyway) and
    # flag it loudly if an out-of-range entry is ever needed by a valid position.
    import deep_gemm

    orig_paged = deep_gemm.fp8_fp4_paged_mqa_logits

    def fp8_fp4_paged_mqa_logits(q, kv_cache, weights, context_lens, block_table, *a, **k):
        num_pages, page_kv = kv_cache.shape[0], kv_cache.shape[1]
        bad = (block_table < 0) | (block_table >= num_pages)
        if bool(bad.any()):
            lens = context_lens.reshape(-1).to(torch.int64)
            needed_pages = (lens + page_kv - 1) // page_kv
            col = torch.arange(block_table.shape[1], device=block_table.device)[None, :]
            if bool((bad & (col < needed_pages[:, None])).any()):
                WORKAROUNDS.append(f"BUG-4b: out-of-range page id inside the valid context: table={block_table.tolist()} lens={lens.tolist()}")
            if "BUG-4" not in WORKAROUNDS:
                WORKAROUNDS.append("BUG-4")
            block_table = torch.where(bad, torch.zeros_like(block_table), block_table)
        return orig_paged(q, kv_cache, weights, context_lens, block_table, *a, **k)

    deep_gemm.fp8_fp4_paged_mqa_logits = fp8_fp4_paged_mqa_logits

    # NOTE-3 (not a bug, a CPU-mode hazard): WorkerBase.profile_num_available_blocks runs a dummy prefill of
    # max_num_batched_tokens tokens (903 s on the host for 2048 tokens) and then sizes the KV cache from
    # mem_get_info(), which on cpu is host RAM -> 0.5 * 2 TB = 559 GiB of KV blocks. Use a fixed small pool instead.
    if num_kv_blocks > 0:
        import torch.distributed as _dist  # noqa: F401
        from deepinfer.worker.base import WorkerBase

        def profile_num_available_blocks(self):
            self.model_runner.set_block_size(self.cache_config.block_size)
            self.cache_config.num_gpu_blocks = num_kv_blocks
            WORKAROUNDS.append(f"NOTE-3: fixed {num_kv_blocks} KV blocks, no profile run")
            return num_kv_blocks

        WorkerBase.profile_num_available_blocks = profile_num_available_blocks


class Recorder:
    """Collects the lm-head logits of every sampler call (one request in flight -> one row per call)."""

    def __init__(self):
        self.step_kind = None  # (is_prompt, is_placeholder) of the engine step being executed
        self.n_placeholder = 0
        self.logits = []
        self.hidden_calls = []
        self._cur_hidden = {}

    def reset(self):
        self.n_placeholder = 0
        self.logits.clear()
        self.hidden_calls.clear()
        self._cur_hidden = {}


def install_logits_hook(rec: Recorder) -> None:
    import deepinfer.model_executor.layers.sampler as sampler_mod

    orig = sampler_mod.get_logits

    def get_logits(hidden_states, embedding, vocab_size):
        logits = orig(hidden_states, embedding, vocab_size)
        if rec.step_kind is not None and rec.step_kind[1]:
            rec.n_placeholder += 1  # placeholder step (no real sequence): not part of the request
            return logits
        rec.logits.append(logits.detach().float().cpu().clone())
        if os.environ.get("RUN_TINY_DEBUG"):
            import traceback as _tb

            st = " <- ".join(f"{f.name}:{f.lineno}" for f in _tb.extract_stack()[-6:-1][::-1])
            print(f"[run]   sampler call {len(rec.logits)}: logits {tuple(logits.shape)} argmax {int(logits[0].argmax())} "
                  f"hs {tuple(hidden_states.shape)} at {time.strftime('%H:%M:%S')} via {st}", flush=True)
        return logits

    sampler_mod.get_logits = get_logits

    from deepinfer.worker.worker import Worker

    orig_step = Worker.execute_model_step

    def execute_model_step(self, is_prompt, is_placeholder, *a, **k):
        rec.step_kind = (bool(is_prompt), bool(is_placeholder))
        try:
            return orig_step(self, is_prompt, is_placeholder, *a, **k)
        finally:
            rec.step_kind = None

    Worker.execute_model_step = execute_model_step


def install_hidden_hooks(rec: Recorder, llm) -> str:
    """Hook the decoder layers of the driver worker's model. Returns a status string."""
    try:
        worker = llm.llm_engine.driver_worker
        runner = getattr(worker, "model_runner", None)
        model = getattr(runner, "model", None) if runner is not None else getattr(worker, "model", None)
        layers = model.model.layers
    except Exception as exc:  # noqa: BLE001
        return f"hidden hooks not installed: {type(exc).__name__}: {exc}"

    from deepinfer.model_executor.models.v4_modules.hc import HCLayerMixin

    n = len(layers)

    def hook(i):
        def fn(mod, inp, out):
            if rec.step_kind is None or rec.step_kind[1]:
                return None  # placeholder / non-request forward
            h, residual = out
            h = h.move() if hasattr(h, "move") else h
            try:
                # 2D output = before hc_post; the reference block output is the expanded hc stream
                x = HCLayerMixin.hc_post(h, residual) if h.dim() == 2 else h
            except Exception:  # noqa: BLE001
                x = h
            rec._cur_hidden[i] = x.detach().float().cpu().clone()
            if i == n - 1:
                rec.hidden_calls.append(dict(rec._cur_hidden))
                rec._cur_hidden = {}
            return None

        return fn

    for i, layer in enumerate(layers):
        layer.register_forward_hook(hook(i))
    return f"hidden hooks on {n} layers of {type(model).__name__}"


def main():
    args = parse_args()
    for k in ("hf_ckpt", "golden_dir", "out"):
        setattr(args, k, os.path.abspath(os.path.expanduser(getattr(args, k))))
    setup_env(args.device)
    os.makedirs(args.out, exist_ok=True)

    import torch

    torch.set_num_threads(args.threads)
    import deepinfer  # noqa: F401
    from deepinfer import PLATFORM

    print(f"[run] platform device = {PLATFORM.device}", flush=True)
    if args.device == "cpu":
        assert str(PLATFORM.device) == "cpu", "refusing to run: target device is not the host CPU"

    install_workarounds(args.num_kv_blocks)
    rec = Recorder()
    install_logits_hook(rec)

    from deepinfer.entrypoints.llm import LLM
    from deepinfer.sampling_params import SamplingParams

    t0 = time.time()
    llm = LLM(
        model=args.hf_ckpt,
        trust_remote_code=True,
        expert_parallel_size=1,
        context_parallel_size=1,
        data_parallel_size=1,
        engine_rank=0,
        chunked_prefill=True,
        decode_sequence=False,
        max_model_len=args.max_model_len,
        max_num_batched_tokens=args.max_num_batched_tokens,
        max_num_seqs=1,
        gpu_memory_utilization=0.5,
    )
    print(f"[run] engine up in {time.time() - t0:.1f}s", flush=True)
    status = install_hidden_hooks(rec, llm) if args.hidden else "hidden not requested"
    print(f"[run] {status}", flush=True)

    for name in args.prompts.split(","):
        gold = torch.load(os.path.join(args.golden_dir, f"{name}.pt"), weights_only=False)
        ids, gtok = list(gold["prompt_ids"]), list(gold["tokens"])
        steps = len(gtok)
        rec.reset()
        sp = SamplingParams(temperature=0.0, max_tokens=steps, ignore_eos=True, golden_token_ids=gtok)
        t1 = time.time()
        outputs = llm.generate(prompt_token_ids=[ids], sampling_params=sp, use_tqdm=False)
        dt = time.time() - t1
        out_tokens = [o.token_id for o in outputs[0]]
        rows = [r for r in rec.logits]
        assert all(r.size(0) == 1 for r in rows), [tuple(r.shape) for r in rows]
        logits = torch.cat(rows)[-steps:] if len(rows) >= steps else torch.cat(rows)
        print(f"[run] {name}: {len(rows)} request sampler calls, {rec.n_placeholder} placeholder steps", flush=True)
        if len(rows) != steps:
            print(f"[run] WARNING {name}: {len(rows)} sampler calls for {steps} steps (chunked prefill?)", flush=True)
        payload = {
            "name": name,
            "prompt_ids": ids,
            "tokens": logits.argmax(-1).tolist(),
            "engine_output_tokens": out_tokens,
            "fed_tokens": gtok[: steps - 1],
            "logits": logits,
            "source": "deepinfer LLM engine, teacher-forced via SamplingParams.golden_token_ids",
            "device": args.device,
        }
        if rec.hidden_calls:
            hid = []
            for call in rec.hidden_calls[-steps:]:
                hid.append(torch.stack([call[i].reshape(call[i].shape[0], -1, gold["hidden"][0].shape[-1]) for i in sorted(call)]).bfloat16())
            payload["hidden"] = hid
        path = os.path.join(args.out, f"{name}.pt")
        torch.save(payload, path)
        match = sum(a == b for a, b in zip(payload["tokens"], gtok))
        print(f"[run] {name}: {len(ids)} prompt tokens, engine output {out_tokens} (golden-forced: {out_tokens == gtok}), "
              f"argmax==golden {match}/{steps}, {dt:.1f}s -> {path}", flush=True)
        if args.compare:
            import compare_golden

            res = compare_golden.compare(gold, payload)
            for s in res["steps"]:
                print(f"    step {s['step']}: match {s['match']} margin {s['margin_ref']:.3f} rel {s['rel']:.2e} cos {s['cos']:.6f} kl {s['kl']:.2e}")
            for L in res["layers"][:1]:
                print("    prefill hidden rel per layer:", L.get("error") or " ".join(f"{v:.1e}" for v in L["rel"]))
    print(f"[run] workarounds triggered: {WORKAROUNDS}", flush=True)
    llm.close()


if __name__ == "__main__":
    try:
        main()
    except Exception:
        traceback.print_exc()
        sys.exit(1)