from __future__ import annotations

import torch


class Animate2KVCacheCPUOffloader:
    """Keep Wan-Animate-2 reference K/V caches in CPU memory.

    The upstream source stores one immutable pre-RoPE K and V tensor per
    transformer layer.  `forward_gen` only reads those tensors.  This allows us
    to preserve one CPU master copy after `forward_ref`, load only the
    current layer's K/V to CUDA immediately before its generation block, and
    restore the dictionary entry to the CPU master immediately after the block.

    This does not quantize or otherwise alter K/V; it is a placement change only.
    """

    def __init__(
        self, model: torch.nn.Module, device: str | torch.device = "cuda", *, pin_memory: bool = False
    ):
        self.model = model
        self.device = torch.device(device)
        self.pin_memory = bool(pin_memory)
        self.handles = []
        self._masters: dict[tuple[int, int, str], torch.Tensor] = {}
        self._current_cache_pair: tuple[int, int] | None = None
        self.ref_offload_bytes = 0
        self.gen_load_bytes = 0
        self.gen_loads = 0
        self.layers_offloaded = set()
        self.peak_layer_bytes = 0

    @staticmethod
    def _nbytes(t: torch.Tensor) -> int:
        return int(t.numel() * t.element_size())

    def _cpu_copy(self, t: torch.Tensor) -> torch.Tensor:
        if t.device.type == "cpu" and (not self.pin_memory or t.is_pinned()):
            return t
        # Keeping an entire 40-layer 480p/720p KV cache page-locked can itself
        # exhaust host pinned memory.  Default to ordinary CPU memory; an
        # explicit cpu-pinned mode is available on hosts with ample RAM.
        out = torch.empty_like(t, device="cpu", pin_memory=self.pin_memory)
        out.copy_(t, non_blocking=False)
        return out

    def _begin_cache_pair(self, k_cache: dict, v_cache: dict, index: int):
        pair = (id(k_cache), id(v_cache))
        # Every segment creates fresh dicts.  When layer 0 of a new reference
        # pass arrives, release the previous segment's pinned masters.
        if int(index) == 0 and self._current_cache_pair != pair:
            self._masters.clear()
            self._current_cache_pair = pair

    def _pre(self, module, args, kwargs):
        if kwargs.get("method") != "forward_gen" or len(args) < 4:
            return
        _x, index, k_cache, v_cache = args[:4]
        index = int(index)
        pair = (id(k_cache), id(v_cache))
        for kind, cache in (("k", k_cache), ("v", v_cache)):
            key = (pair[0] if kind == "k" else pair[1], index, kind)
            master = self._masters.get(key)
            if master is None:
                # This indicates the caller bypassed the reference pass or the
                # cache dictionaries changed unexpectedly.  Do not silently
                # fall back to GPU residency.
                raise RuntimeError(
                    f"missing CPU {kind.upper()} master for Animate-2 layer {index}; "
                    "forward_ref must complete before forward_gen"
                )
            gpu = master.to(self.device, non_blocking=self.pin_memory)
            cache[index] = gpu
            self.gen_load_bytes += self._nbytes(master)
        self.gen_loads += 1

    def _post(self, module, args, kwargs, output):
        method = kwargs.get("method")
        if method not in {"forward_ref", "forward_gen"} or len(args) < 4:
            return output
        _x, index, k_cache, v_cache = args[:4]
        index = int(index)
        pair = (id(k_cache), id(v_cache))

        if method == "forward_ref":
            self._begin_cache_pair(k_cache, v_cache, index)
            layer_bytes = 0
            for kind, cache in (("k", k_cache), ("v", v_cache)):
                t = cache.get(index)
                if t is None:
                    raise RuntimeError(f"Animate-2 reference pass did not populate {kind}_cache[{index}]")
                master = self._cpu_copy(t.detach())
                key = (pair[0] if kind == "k" else pair[1], index, kind)
                self._masters[key] = master
                cache[index] = master
                nb = self._nbytes(master)
                self.ref_offload_bytes += nb
                layer_bytes += nb
            self.layers_offloaded.add(index)
            self.peak_layer_bytes = max(self.peak_layer_bytes, layer_bytes)
        else:
            # K/V are immutable in forward_gen.  Restore the already-existing
            # pinned master instead of copying the CUDA tensor back to host.
            for kind, cache in (("k", k_cache), ("v", v_cache)):
                key = (pair[0] if kind == "k" else pair[1], index, kind)
                master = self._masters.get(key)
                if master is None:
                    raise RuntimeError(f"lost CPU {kind.upper()} master for layer {index}")
                cache[index] = master
        return output

    def install(self):
        if self.device.type != "cuda":
            raise ValueError("Animate2KVCacheCPUOffloader requires a CUDA device")
        for block in self.model.blocks:
            self.handles.append(block.register_forward_pre_hook(self._pre, with_kwargs=True))
            self.handles.append(block.register_forward_hook(self._post, with_kwargs=True))
        return self

    def remove(self):
        for handle in self.handles:
            handle.remove()
        self.handles.clear()
        self._masters.clear()

    def report(self) -> dict:
        return {
            "mode": "pinned_cpu_per_layer_stream" if self.pin_memory else "cpu_per_layer_stream",
            "pinned_host_master": self.pin_memory,
            "layers_offloaded": len(self.layers_offloaded),
            "reference_offload_gib": self.ref_offload_bytes / 2**30,
            "generation_h2d_gib": self.gen_load_bytes / 2**30,
            "generation_layer_loads": int(self.gen_loads),
            "peak_single_layer_cache_gib": self.peak_layer_bytes / 2**30,
            "full_cache_gpu_residency": False,
        }