"""Opt-in tensor dumps for H20 bring-up comparisons.

Enabled with `DEEPINFER_DUMP_LAYERS=<dir>`. Each call appends one row to
`<dir>/deepinfer_layers.npz`-style storage under a per-tag key, always as the *last* row of
the batch: the reference implementation only generates from the final position, so that is
the only row the two sides can be compared on.

The dump is deliberately simple (host copy + `numpy.savez` at exit) because it exists to
answer "which layer diverges first", not to be a profiler.
"""

from __future__ import annotations

import os
from pathlib import Path

import numpy as np
import torch

_DUMP_DIR = os.getenv("DEEPINFER_DUMP_LAYERS")
# The engine runs a dummy profile forward before it allocates the KV cache, and that one has a
# different token count. Gating on the real prompt length keeps the dump comparable.
_DUMP_TOKENS = int(os.getenv("DEEPINFER_DUMP_TOKENS", "0") or 0)
_CAPTURED: dict[str, np.ndarray] = {}
# Only the first N layers of MoE internals are kept: the routing and per-branch outputs are the
# interesting part for bring-up, and dumping all 40 layers would slow the forward down for no gain.
_MOE_LAYERS = int(os.getenv("DEEPINFER_DUMP_MOE_LAYERS", "2") or 0)
# Per-tag call counter for `record_step`.
_STEPS: dict[str, int] = {}
# Number of rows in the most recent logits tensor, used to skip warmup passes.
_ROWS = 0
# A decode step that serves real requests has at most this many rows. Anything wider is warmup.
_REAL_ROWS = int(os.getenv("DEEPINFER_DUMP_MAX_ROWS", "1") or 1)


_SUSPENDED = False


def enabled() -> bool:
    return bool(_DUMP_DIR) and not _SUSPENDED


class profiling:
    """Suspend capture while the engine runs its dummy profile forward.

    `execute_model_profile_run` goes through the same forward and sampler as a real step, with
    synthetic tokens, so anything captured there is garbage that later masking cannot detect. The
    layer dumps happen to be filtered by their token-count gate; the logits are not, because their
    row count is the number of *selected* tokens rather than the prompt length.
    """

    def __enter__(self) -> None:
        global _SUSPENDED
        self._previous = _SUSPENDED
        _SUSPENDED = True

    def __exit__(self, *_exc: object) -> None:
        global _SUSPENDED
        _SUSPENDED = self._previous


def record(tag: str, tensor: torch.Tensor, once: bool = False) -> None:
    """Store the last row of `tensor` under `tag`.

    A 3D tensor is treated as `[tokens, hc_mult, dim]` and reduced with `mean(dim=1)`, which is
    the closest match to the reference's collapsed residual stream.

    `once=True` keeps the first write only. Logits need this: the row count is the number of
    selected tokens rather than the prompt length (so the token gate does not apply), and the
    decode step that follows prefill would otherwise overwrite the value we care about.
    """
    if not _DUMP_DIR:
        return
    if _DUMP_TOKENS and not once and tensor.shape[0] != _DUMP_TOKENS:
        return
    if once and tag in _CAPTURED:
        return
    if not _want_moe(tag):
        return
    if tensor.dim() == 3:
        tensor = tensor.mean(dim=1)
    row = tensor[-1] if tensor.dim() == 2 else tensor.reshape(-1)
    # Copying to the host is the whole point here, and it happens inside the forward where the
    # engine forbids syncs, so take the sanctioned escape hatch for this one transfer.
    from deepinfer.utils import allow_cuda_sync

    with allow_cuda_sync():
        _CAPTURED[tag] = row.detach().float().cpu().numpy()


def set_rows(rows: int) -> None:
    """Tell the dump how many rows the sampler just produced, so warmup can be filtered out."""
    global _ROWS
    _ROWS = rows


def record_step(tag: str, tensor: torch.Tensor) -> None:
    """Append `tensor`'s last row under `tag_<n>`, numbering the calls.

    Decode needs every step kept, not just the first: the point of the decode comparison is to
    watch the margin evolve and find the step where it collapses.
    """
    if not _DUMP_DIR:
        return
    # Decode with one request produces exactly `_REAL_ROWS` rows; warmup batches are larger.
    if _ROWS > _REAL_ROWS:
        return
    index = _STEPS.get(tag, 0)
    _STEPS[tag] = index + 1
    # Row 0, not the last row: during decode the batch is padded (CUDA graph slots, or just
    # `max_num_seqs`), and the padding rows hold logits for placeholder tokens. Row 0 is the
    # first real request's current token, which is the one to compare against a batch-1 reference.
    row = tensor[0] if tensor.dim() >= 2 else tensor
    from deepinfer.utils import allow_cuda_sync

    with allow_cuda_sync():
        _CAPTURED[f"{tag}_{index}"] = row.detach().float().cpu().numpy()


def record_ids(tag: str, tensor: torch.Tensor) -> None:
    """Store an integer tensor as-is (expert ids); reduction would destroy its meaning."""
    if not _DUMP_DIR:
        return
    if not _want_moe(tag):
        return
    if _DUMP_TOKENS and tensor.shape[0] != _DUMP_TOKENS:
        return
    from deepinfer.utils import allow_cuda_sync

    with allow_cuda_sync():
        _CAPTURED[tag] = tensor.detach().to(torch.int32).cpu().numpy()


def _want_moe(tag: str) -> bool:
    """Keep MoE internals only for the first `_MOE_LAYERS` layers."""
    if "_" not in tag:
        return True
    head = tag.split("_", 1)[0]
    if not head.startswith("layer"):
        return True
    try:
        return int(head[5:]) < _MOE_LAYERS
    except ValueError:
        return True


def _rank_suffix() -> str:
    """`_rank<N>` when running distributed, so ranks do not overwrite each other's dump.

    With data parallelism every rank registers `atexit(flush)` and writes the same filename; the
    last one to exit wins, which silently mixes one rank's tensors with another's. Only rank 0
    holds the real prompt anyway, so the suffix makes that explicit instead of accidental.
    """
    rank = os.getenv("RANK")
    if rank is None:
        try:
            import torch.distributed as dist

            if dist.is_initialized():
                rank = str(dist.get_rank())
        except Exception:  # noqa: BLE001 - torch may not be usable at exit
            rank = None
    # With data parallelism each engine is its own single-process group, so every rank reports
    # RANK=0 and they all target the same file. The pid disambiguates them.
    return f"_rank{rank if rank is not None else 'x'}_pid{os.getpid()}"


def flush() -> None:
    """Write everything captured so far. Called once at interpreter exit."""
    if not _DUMP_DIR or not _CAPTURED:
        return
    out = Path(_DUMP_DIR)
    out.mkdir(parents=True, exist_ok=True)
    path = out / f"deepinfer_layers{_rank_suffix()}.npz"
    np.savez(path, **_CAPTURED)
    print(f"[dump] wrote {len(_CAPTURED)} tensors to {path}", flush=True)
    _CAPTURED.clear()


if _DUMP_DIR:
    import atexit

    atexit.register(flush)