"""
Hybrid CPU Offload

Runs a model across GPU + CPU: keeps embedding, final-norm, lm_head,
and a sliding window of transformer layers on GPU. Remaining layers
live on CPU (pinned memory) and are streamed to GPU on demand with
async prefetch.

Architecture:
    1.  Model loads to CPU (low memory).
    2.  Static components (embed, norm, head, rotary_emb) → GPU.
    3.  LayerStreamer manages a GPU window of N layers.
    4.  During forward, layers are pulled to GPU one-by-one;
        the prefetcher brings the next few layers while compute
        runs on the current layer.
    5.  After use each layer is marked evictable; when the window
        is full the LRU layer is offloaded back to CPU (pinned).

This lets a 7 B FP16 model (~14 GB) run on a GPU with only 6-8 GB
free VRAM, at the cost of some PCIe transfer latency that the
prefetcher mostly hides.
"""

from __future__ import annotations

import time
import threading
from collections import OrderedDict
from contextlib import nullcontext
from dataclasses import dataclass
from typing import Any, Dict, List, Optional, Tuple, Union

import torch
import torch.nn as nn


# ---------------------------------------------------------------------------
# Config
# ---------------------------------------------------------------------------

@dataclass
class OffloadConfig:
    """Configuration for hybrid CPU offload."""
    gpu_layers: int = 4          # Max transformer layers resident on GPU
    prefetch_count: int = 2      # How many next-layers to prefetch
    use_pinned_memory: bool = True
    device: int = 0              # CUDA device ordinal


# ---------------------------------------------------------------------------
# Layer Streaming Engine (lightweight, self-contained)
# ---------------------------------------------------------------------------

class _LayerWindow:
    """
    Manages a sliding window of transformer layers on GPU.

    Layers outside the window live on CPU (pinned memory when possible).
    Access triggers load-to-GPU; eviction is LRU.
    """

    def __init__(
        self,
        layers: nn.ModuleList,
        max_window: int,
        device: int = 0,
        use_pinned: bool = True,
        prefetch_count: int = 2,
    ):
        self.layers = layers
        self.num_layers = len(layers)
        self.max_window = min(max_window, self.num_layers)
        self.device = device
        self.gpu_device = f"cuda:{device}"
        self.use_pinned = use_pinned
        self.prefetch_count = prefetch_count

        # LRU window: OrderedDict[layer_idx] -> True
        self._window: OrderedDict[int, bool] = OrderedDict()
        self._lock = threading.Lock()

        # Separate CUDA stream for async H2D transfers
        self._transfer_stream = torch.cuda.Stream(device=device)

        # Stats
        self.stats = {
            "gpu_hits": 0,
            "cpu_loads": 0,
            "evictions": 0,
            "total_transfer_s": 0.0,
        }

    # -- public API --

    def ensure_on_gpu(self, layer_idx: int) -> nn.Module:
        """Return the layer Module on GPU, loading it if necessary."""
        with self._lock:
            if layer_idx in self._window:
                self._window.move_to_end(layer_idx)
                self.stats["gpu_hits"] += 1
                return self.layers[layer_idx]

            # Need to load → make room first
            while len(self._window) >= self.max_window:
                self._evict_lru()

            self._load_to_gpu(layer_idx)
            return self.layers[layer_idx]

    def prefetch(self, current_idx: int) -> None:
        """Async-prefetch the next *prefetch_count* layers after *current_idx*."""
        for offset in range(1, self.prefetch_count + 1):
            nxt = current_idx + offset
            if nxt >= self.num_layers:
                break
            with self._lock:
                if nxt in self._window:
                    continue
                if len(self._window) >= self.max_window:
                    self._evict_lru()
                self._load_to_gpu(nxt, async_transfer=True)

    def release(self, layer_idx: int) -> None:
        """Mark layer as finished (currently a no-op; stays until evicted)."""
        pass

    # -- internals --

    @staticmethod
    def _deep_move(module: nn.Module, device: str, non_blocking: bool = True) -> None:
        """
        Move a module to *device*, including non-registered tensor attrs
        (e.g. QuantizedLinearZSE._cached_weight, _triton_v2_weight, etc.).
        """
        # Standard .to() handles parameters + registered buffers
        module.to(device, non_blocking=non_blocking)

        # Also move any plain-attribute tensors on every submodule
        _EXTRA_ATTRS = (
            "_cached_weight",
            "_triton_v2_weight",
            "_triton_v2_scales",
            "_bnb_weight",
        )
        for sub in module.modules():
            for attr in _EXTRA_ATTRS:
                t = getattr(sub, attr, None)
                if isinstance(t, torch.Tensor):
                    setattr(sub, attr, t.to(device, non_blocking=non_blocking))

    @staticmethod
    def _deep_cpu(module: nn.Module, pin: bool = True) -> None:
        """
        Move a module to CPU, including non-registered tensor attrs.
        Optionally pin memory for faster future H2D transfers.
        """
        _EXTRA_ATTRS = (
            "_cached_weight",
            "_triton_v2_weight",
            "_triton_v2_scales",
            "_bnb_weight",
        )
        # Move params
        for p in module.parameters():
            p.data = p.data.cpu().pin_memory() if pin else p.data.cpu()
        # Move registered buffers
        for _, buf in module.named_buffers():
            if buf is not None:
                buf.data = buf.data.cpu().pin_memory() if pin else buf.data.cpu()
        # Move non-registered tensor attrs
        for sub in module.modules():
            for attr in _EXTRA_ATTRS:
                t = getattr(sub, attr, None)
                if isinstance(t, torch.Tensor):
                    moved = t.cpu()
                    setattr(sub, attr, moved.pin_memory() if pin else moved)

    def _load_to_gpu(self, idx: int, async_transfer: bool = False) -> None:
        t0 = time.perf_counter()
        layer = self.layers[idx]
        ctx = torch.cuda.stream(self._transfer_stream) if async_transfer else nullcontext()
        with ctx:
            self._deep_move(layer, self.gpu_device, non_blocking=True)
        if not async_transfer:
            torch.cuda.synchronize(self.device)
        self._window[idx] = True
        self.stats["cpu_loads"] += 1
        self.stats["total_transfer_s"] += time.perf_counter() - t0

    def _evict_lru(self) -> None:
        if not self._window:
            return
        idx, _ = self._window.popitem(last=False)
        self._deep_cpu(self.layers[idx], pin=self.use_pinned)
        self.stats["evictions"] += 1


# ---------------------------------------------------------------------------
# Offload Model Wrapper  (nn.Module compatible with TextGenerator)
# ---------------------------------------------------------------------------

@dataclass
class OffloadModelOutput:
    """Mimics transformers CausalLMOutput."""
    logits: torch.Tensor
    past_key_values: Any = None


class OffloadModelWrapper(nn.Module):
    """
    Wraps a HuggingFace causal-LM loaded on **CPU** so that it runs
    with hybrid GPU + CPU layer streaming.

    Static tensors (embed, norm, head, rotary_emb) live on GPU.
    Transformer layers are streamed through a GPU window.
    """

    def __init__(self, model: nn.Module, config: OffloadConfig):
        super().__init__()
        self.inner = model
        self.cfg = config
        self._device = torch.device(f"cuda:{config.device}")

        # Detect arch (Qwen2/Llama share "model.layers"; GPT-2 uses "transformer.h")
        self.arch, self._layers_attr = self._detect_arch()
        layers = self._get_layers()
        self.num_layers = len(layers)

        # Move static components to GPU
        self._pin_static_to_gpu()

        # Move all layers to CPU (pinned)
        for layer in layers:
            _LayerWindow._deep_cpu(layer, pin=config.use_pinned_memory)

        # Layer window (streaming engine)
        self._window = _LayerWindow(
            layers=layers,
            max_window=config.gpu_layers,
            device=config.device,
            use_pinned=config.use_pinned_memory,
            prefetch_count=config.prefetch_count,
        )

        # Cache info
        self._kv_cache: Any = None

        print(
            f"[HybridOffload] {self.arch} · {self.num_layers} layers · "
            f"GPU window={config.gpu_layers} · prefetch={config.prefetch_count}"
        )

    # ------------------------------------------------------------------
    # Architecture detection
    # ------------------------------------------------------------------

    _ARCH_MAP = [
        # (attr_path_to_layers, arch_name)
        ("model.layers", "llama"),        # Llama / Qwen2 / Mistral
        ("transformer.h", "gpt2"),        # GPT-2
        ("gpt_neox.layers", "gpt_neox"),  # GPT-NeoX
        ("transformer.blocks", "falcon"), # Falcon
    ]

    def _detect_arch(self) -> Tuple[str, str]:
        for attr, arch in self._ARCH_MAP:
            try:
                obj = self.inner
                for part in attr.split("."):
                    obj = getattr(obj, part)
                if isinstance(obj, nn.ModuleList) and len(obj) > 0:
                    return arch, attr
            except AttributeError:
                continue
        raise ValueError("Cannot detect model architecture for offload")

    def _get_layers(self) -> nn.ModuleList:
        obj = self.inner
        for part in self._layers_attr.split("."):
            obj = getattr(obj, part)
        return obj

    # ------------------------------------------------------------------
    # Static component pinning
    # ------------------------------------------------------------------

    def _pin_static_to_gpu(self) -> None:
        """Move embed, norm, head, rotary_emb to GPU."""
        dev = str(self._device)
        m = self.inner

        if self.arch == "llama":
            m.model.embed_tokens.to(dev)
            m.model.norm.to(dev)
            m.lm_head.to(dev)
            if hasattr(m.model, "rotary_emb"):
                # Try .to() first; if it fails (meta tensor), recreate
                try:
                    m.model.rotary_emb.to(dev)
                except Exception:
                    from transformers.models.qwen2.modeling_qwen2 import Qwen2RotaryEmbedding
                    m.model.rotary_emb = Qwen2RotaryEmbedding(config=m.config).to(dev)
        elif self.arch == "gpt2":
            m.transformer.wte.to(dev)
            if hasattr(m.transformer, "wpe"):
                m.transformer.wpe.to(dev)
            m.transformer.ln_f.to(dev)
            m.lm_head.to(dev)
        elif self.arch == "gpt_neox":
            m.gpt_neox.embed_in.to(dev)
            m.gpt_neox.final_layer_norm.to(dev)
            m.embed_out.to(dev)
        elif self.arch == "falcon":
            m.transformer.word_embeddings.to(dev)
            m.transformer.ln_f.to(dev)
            m.lm_head.to(dev)

    # ------------------------------------------------------------------
    # Forward
    # ------------------------------------------------------------------

    def forward(
        self,
        input_ids: Optional[torch.Tensor] = None,
        attention_mask: Optional[torch.Tensor] = None,
        position_ids: Optional[torch.Tensor] = None,
        past_key_values: Any = None,
        use_cache: bool = True,
        **kwargs,
    ) -> OffloadModelOutput:
        dev = self._device
        m = self.inner

        # --- Embeddings ---
        if self.arch == "llama":
            hidden = m.model.embed_tokens(input_ids.to(dev))
        elif self.arch == "gpt2":
            hidden = m.transformer.wte(input_ids.to(dev))
            if hasattr(m.transformer, "wpe"):
                seq_len = input_ids.shape[1]
                past_len = 0
                if past_key_values is not None and len(past_key_values) > 0:
                    first = past_key_values[0]
                    if first is not None:
                        past_len = first[0].shape[-2] if isinstance(first, tuple) else first.key_cache[0].shape[-2]
                pos = torch.arange(past_len, past_len + seq_len, device=dev)
                hidden = hidden + m.transformer.wpe(pos)
        else:
            raise NotImplementedError(f"Offload forward not implemented for {self.arch}")

        # --- Position IDs / embeddings ---
        seq_len = hidden.shape[1]
        if position_ids is None:
            past_len = 0
            if past_key_values is not None:
                if hasattr(past_key_values, "get_seq_length"):
                    past_len = past_key_values.get_seq_length()
                elif isinstance(past_key_values, (list, tuple)) and len(past_key_values) > 0:
                    first = past_key_values[0]
                    if first is not None:
                        past_len = first[0].shape[-2] if isinstance(first, tuple) else 0
            position_ids = torch.arange(
                past_len, past_len + seq_len, dtype=torch.long, device=dev
            ).unsqueeze(0)

        # Compute rotary embeddings once (Qwen2 / Llama)
        position_embeddings = None
        if self.arch == "llama" and hasattr(m.model, "rotary_emb"):
            position_embeddings = m.model.rotary_emb(hidden, position_ids)

        # Prepare DynamicCache for transformers >= 4.57
        if use_cache and past_key_values is None:
            from transformers import DynamicCache
            past_key_values = DynamicCache()

        # --- Layer-by-layer streaming ---
        for i in range(self.num_layers):
            # Bring layer to GPU
            layer = self._window.ensure_on_gpu(i)

            # Kick off prefetch for next layers
            self._window.prefetch(i)

            # Forward through layer
            if self.arch == "llama":
                layer_out = layer(
                    hidden,
                    attention_mask=attention_mask,
                    position_ids=position_ids,
                    past_key_value=past_key_values,
                    use_cache=use_cache,
                    position_embeddings=position_embeddings,
                )
            elif self.arch == "gpt2":
                layer_out = layer(
                    hidden,
                    attention_mask=attention_mask,
                    layer_past=past_key_values[i] if isinstance(past_key_values, list) and i < len(past_key_values) else None,
                    use_cache=use_cache,
                )
            else:
                layer_out = layer(hidden)

            hidden = layer_out[0]
            self._window.release(i)

        # --- Final norm + head ---
        if self.arch == "llama":
            hidden = m.model.norm(hidden)
            logits = m.lm_head(hidden)
        elif self.arch == "gpt2":
            hidden = m.transformer.ln_f(hidden)
            logits = m.lm_head(hidden)
        else:
            logits = hidden  # fallback

        return OffloadModelOutput(
            logits=logits,
            past_key_values=past_key_values if use_cache else None,
        )

    # ------------------------------------------------------------------
    # generate() – simple autoregressive loop
    # ------------------------------------------------------------------

    def generate(
        self,
        input_ids: torch.Tensor,
        max_new_tokens: int = 128,
        temperature: float = 0.7,
        top_p: float = 0.9,
        **kwargs,
    ) -> torch.Tensor:
        dev = self._device
        ids = input_ids.to(dev)
        past = None

        for _ in range(max_new_tokens):
            inp = ids[:, -1:] if past is not None else ids
            out = self.forward(input_ids=inp, past_key_values=past, use_cache=True)
            past = out.past_key_values

            logits = out.logits[:, -1, :] / max(temperature, 1e-7)

            # top-p
            if top_p < 1.0:
                sorted_logits, sorted_idx = torch.sort(logits, descending=True)
                cum = torch.cumsum(torch.softmax(sorted_logits, dim=-1), dim=-1)
                remove = cum > top_p
                remove[:, 1:] = remove[:, :-1].clone()
                remove[:, 0] = False
                logits[sorted_idx[remove]] = float("-inf")

            probs = torch.softmax(logits, dim=-1)
            nxt = torch.multinomial(probs, 1) if temperature > 0 else logits.argmax(-1, keepdim=True)
            ids = torch.cat([ids, nxt], dim=-1)

            # EOS
            if hasattr(self.inner, "config") and hasattr(self.inner.config, "eos_token_id"):
                eos = self.inner.config.eos_token_id
                if isinstance(eos, int) and nxt.item() == eos:
                    break
                if isinstance(eos, list) and nxt.item() in eos:
                    break

        return ids

    # ------------------------------------------------------------------
    # Compatibility helpers
    # ------------------------------------------------------------------

    def parameters(self):
        return self.inner.parameters()

    def eval(self):
        self.inner.eval()
        return self

    def to(self, *args, **kwargs):
        # Offload wrapper manages placement; ignore external .to() calls
        return self

    @property
    def device(self):
        return self._device

    @property
    def config(self):
        return getattr(self.inner, "config", None)

    def get_stats(self) -> Dict[str, Any]:
        return {
            "arch": self.arch,
            "num_layers": self.num_layers,
            "gpu_window": self.cfg.gpu_layers,
            **self._window.stats,
        }


# ---------------------------------------------------------------------------
# Public API
# ---------------------------------------------------------------------------

class HybridOffloadCoordinator:
    """
    High-level coordinator that wraps a HuggingFace model for hybrid
    GPU+CPU inference with layer streaming.

    Usage from orchestrator:
        coordinator = HybridOffloadCoordinator(model, device=0, gpu_layers=6)
        wrapper = coordinator.get_wrapper()   # nn.Module
    """

    def __init__(
        self,
        model: nn.Module,
        device: int = 0,
        gpu_layers: int = 4,
        prefetch_count: int = 2,
        use_pinned_memory: bool = True,
    ):
        config = OffloadConfig(
            gpu_layers=gpu_layers,
            prefetch_count=prefetch_count,
            use_pinned_memory=use_pinned_memory,
            device=device,
        )
        self.wrapper = OffloadModelWrapper(model, config)

    def get_wrapper(self) -> OffloadModelWrapper:
        return self.wrapper

    def get_stats(self) -> Dict[str, Any]:
        return self.wrapper.get_stats()

    def shutdown(self) -> None:
        """Cleanup (move everything to CPU to free VRAM)."""
        layers = self.wrapper._get_layers()
        for layer in layers:
            layer.cpu()
        torch.cuda.empty_cache()
