"""
torrent/models/__init__.py
==========================
Qwen3MoE × Torrent 集成

Torrent 注入方式：
  Monkey-patch Qwen3MoeSparseMoeBlock.forward → TorrentMoEForward.__call__
  通过 TorrentState 跨步保存 prev_routing（逐层），驱动热专家预取。

Expert 存储：
  每层 128 个专家权重以 CPU pin_memory 形式存储在 TorrentLayerExpertStore 中，
  通过独立 CUDA io_stream 异步 H2D，计算在 compute_stream 上执行。
"""

from __future__ import annotations

import logging
import time
import types
from typing import TYPE_CHECKING, Dict, List, Optional, Tuple

import torch
import torch.nn as nn
import torch.nn.functional as F

from srep_moe.prefetch.correlation_table import CorrelationTable
from srep_moe.metrics.collector import MetricsCollector

if TYPE_CHECKING:
    from srep_moe.metrics.tracer import ExpertTracer

logger = logging.getLogger(__name__)


# ============================================================================
#  per-layer Expert Cache
# ============================================================================

class TorrentLayerExpertStore:
    """
    单层的专家权重管理 + 异步 H2D 预取。

    CPU pin_memory 常驻所有专家权重；
    GPU 侧只缓存当前层 forward 中正在使用的专家（用完即释放）。
    """

    def __init__(
        self,
        layer_idx: int,
        num_experts: int,
        device: torch.device,
        io_stream: Optional["torch.cuda.Stream"],
        compute_stream: Optional["torch.cuda.Stream"],
        tracer: Optional["ExpertTracer"] = None,
    ):
        self.layer_idx = layer_idx
        self.num_experts = num_experts
        self.device = device
        self.io_stream = io_stream
        self.compute_stream = compute_stream
        self._use_cuda = (device.type == "cuda") and (io_stream is not None)
        self._tracer = tracer

        # 当前推理阶段（由 set_phase() 设置）
        self._current_phase: str = "ar"
        self._current_step: int = 0
        self._current_draft_pos: Optional[int] = None

        self._cpu_weights: Dict[int, Dict[str, torch.Tensor]] = {}
        self._gpu_cache: Dict[int, Dict[str, torch.Tensor]] = {}
        self._in_flight: set = set()
        self._prefetch_times: Dict[int, float] = {}
        self._used_experts: set = set()

    def register(self, eid: int, gate_up_cpu: torch.Tensor, down_cpu: torch.Tensor) -> None:
        if self._use_cuda:
            gu = gate_up_cpu.pin_memory() if not gate_up_cpu.is_pinned() else gate_up_cpu
            dw = down_cpu.pin_memory() if not down_cpu.is_pinned() else down_cpu
        else:
            gu, dw = gate_up_cpu, down_cpu
        self._cpu_weights[eid] = {"gate_up_proj": gu, "down_proj": dw}

    def set_phase(
        self,
        phase: str,
        step: int,
        draft_pos: Optional[int] = None,
    ) -> None:
        """设置当前推理阶段（由 EAGLE3SDDriver 或推理循环在 forward 前调用）。"""
        self._current_phase = phase
        self._current_step = step
        self._current_draft_pos = draft_pos

    def reset(self) -> None:
        self._gpu_cache.clear()
        self._in_flight.clear()
        self._prefetch_times.clear()
        self._used_experts.clear()

    def prefetch(self, expert_ids: List[int], context: str = "hot") -> None:
        if not self._use_cuda:
            for eid in expert_ids:
                if eid not in self._gpu_cache and eid in self._cpu_weights:
                    self._gpu_cache[eid] = {
                        k: v.to(self.device)
                        for k, v in self._cpu_weights[eid].items()
                    }
                    if self._tracer:
                        self._tracer.emit({
                            "event": "prefetch_start",
                            "expert_id": eid,
                            "layer": self.layer_idx,
                            "context": context,
                            "step": self._current_step,
                            "phase": self._current_phase,
                            **({
                                "draft_pos": self._current_draft_pos
                            } if self._current_draft_pos is not None else {}),
                        })
                        self._tracer.emit({
                            "event": "prefetch_ready",
                            "expert_id": eid,
                            "layer": self.layer_idx,
                            "transfer_ms": 0.0,
                        })
            return
        with torch.cuda.stream(self.io_stream):
            for eid in expert_ids:
                if eid not in self._gpu_cache and eid in self._cpu_weights:
                    start_t = time.perf_counter()
                    self._gpu_cache[eid] = {
                        k: v.to(self.device, non_blocking=True)
                        for k, v in self._cpu_weights[eid].items()
                    }
                    self._in_flight.add(eid)
                    self._prefetch_times[eid] = start_t
                    if self._tracer:
                        ev: Dict = {
                            "event": "prefetch_start",
                            "expert_id": eid,
                            "layer": self.layer_idx,
                            "context": context,
                            "step": self._current_step,
                            "phase": self._current_phase,
                        }
                        if self._current_draft_pos is not None:
                            ev["draft_pos"] = self._current_draft_pos
                        self._tracer.emit(ev)

    def wait_and_get(self, eid: int) -> Dict[str, torch.Tensor]:
        if eid not in self._gpu_cache:
            if self._use_cuda:
                with torch.cuda.stream(self.io_stream):
                    self._gpu_cache[eid] = {
                        k: v.to(self.device, non_blocking=True)
                        for k, v in self._cpu_weights[eid].items()
                    }
                self._in_flight.add(eid)
                self._prefetch_times.setdefault(eid, time.perf_counter())
            else:
                self._gpu_cache[eid] = {
                    k: v.to(self.device)
                    for k, v in self._cpu_weights[eid].items()
                }

        in_flight_before = eid in self._in_flight

        if self._use_cuda and eid in self._in_flight:
            self.compute_stream.wait_stream(self.io_stream)
            self._in_flight.discard(eid)

        # 发出 prefetch_ready 事件（只在 H2D 等待完成后才记录）
        if self._tracer and in_flight_before:
            start_t = self._prefetch_times.pop(eid, None)
            transfer_ms = (time.perf_counter() - start_t) * 1e3 if start_t else 0.0
            self._tracer.emit({
                "event": "prefetch_ready",
                "expert_id": eid,
                "layer": self.layer_idx,
                "transfer_ms": round(transfer_ms, 4),
            })

        self._used_experts.add(eid)
        return self._gpu_cache[eid]

    def is_ready(self, eid: int) -> bool:
        if eid not in self._gpu_cache:
            return False
        if eid in self._in_flight:
            return False
        return True

    def evict(self, expert_ids: List[int]) -> None:
        for eid in expert_ids:
            was_used = eid in self._used_experts
            if self._tracer and eid in self._gpu_cache:
                self._tracer.emit({
                    "event": "expert_evict",
                    "expert_id": eid,
                    "layer": self.layer_idx,
                    "was_used": was_used,
                    "step": self._current_step,
                    "phase": self._current_phase,
                })
            self._gpu_cache.pop(eid, None)
            self._in_flight.discard(eid)
            self._prefetch_times.pop(eid, None)
            self._used_experts.discard(eid)


# ============================================================================
#  Torrent 共享状态（跨 step 保存路由历史）
# ============================================================================

class TorrentState:
    """
    跨 step/层的 Torrent 推理状态。

    prev_routing[layer_idx] = 上一个 decode step 在该层的 routing_indices [T, top_k]
    """

    def __init__(self, num_layers: int):
        self.num_layers = num_layers
        self.prev_routing: Dict[int, Optional[torch.Tensor]] = {
            i: None for i in range(num_layers)
        }
        self._current: Dict[int, Optional[torch.Tensor]] = {
            i: None for i in range(num_layers)
        }

    def update_layer(self, layer_idx: int, routing: torch.Tensor) -> None:
        self._current[layer_idx] = routing.detach().cpu()

    def advance(self) -> None:
        self.prev_routing = dict(self._current)
        self._current = {i: None for i in range(self.num_layers)}

    def reset(self) -> None:
        self.prev_routing = {i: None for i in range(self.num_layers)}
        self._current = {i: None for i in range(self.num_layers)}


# ============================================================================
#  单层 Torrent MoE Forward
# ============================================================================

class TorrentMoEForward:
    """
    替换 Qwen3MoeSparseMoeBlock.forward 的 Torrent 版本。

    调用接口（与原始相同）：
        __call__(hidden_states: [B, T, H]) → [B, T, H]
    """

    def __init__(
        self,
        moe_block,
        layer_idx: int,
        store: TorrentLayerExpertStore,
        corr_table: CorrelationTable,
        state: TorrentState,
        metrics: MetricsCollector,
        top_k: int = 8,
        prefetch_k: int = 8,
        device: torch.device = torch.device("cuda"),
        act_fn=F.silu,
        tracer: Optional["ExpertTracer"] = None,
    ):
        self.moe_block = moe_block
        self.layer_idx = layer_idx
        self.store = store
        self.corr = corr_table
        self.state = state
        self.metrics = metrics
        self.top_k = top_k
        self.prefetch_k = prefetch_k
        self.device = device
        self.act_fn = act_fn
        self._tracer = tracer

    def __call__(self, hidden_states: torch.Tensor) -> torch.Tensor:
        batch_size, seq_len, hidden_dim = hidden_states.shape
        hs_2d = hidden_states.view(-1, hidden_dim)  # [T, H]

        # 同步 tracer 的当前 phase/step 到 store（让 prefetch 事件携带正确上下文）
        if self._tracer:
            self.store.set_phase(
                self._tracer.current_phase,
                self._tracer.current_step,
                self._tracer.current_draft_pos,
            )

        # 1. 预测热专家
        prev_r = self.state.prev_routing.get(self.layer_idx)
        hot_ids: List[int] = []
        if prev_r is not None and prev_r.numel() > 0:
            prev_r_dev = prev_r.to(self.device)
            hot_ids = self.corr.predict_hot_experts(
                self.layer_idx, prev_r_dev, self.prefetch_k
            )

        # 2. 预取热专家（与 gate 计算重叠）
        self.store.reset()
        if hot_ids:
            self.store.prefetch(hot_ids, context="hot")

        # 3. Gate 前向
        _, routing_weights, routing_indices = self.moe_block.gate(hs_2d)

        # 4. 确定冷专家，立即启动预取
        active_list = routing_indices.reshape(-1).tolist()
        active_set = set(int(e) for e in active_list)
        hot_set = {int(e) for e in hot_ids} & active_set
        cold_set = active_set - hot_set

        if cold_set:
            self.store.prefetch(list(cold_set), context="cold")

        # 5. 构建 expert→token 映射
        num_tokens = hs_2d.shape[0]
        expert_tokens: Dict[int, Tuple[List[int], List[float]]] = {}
        for ti in range(num_tokens):
            for k in range(self.top_k):
                eid = int(routing_indices[ti, k].item())
                w = float(routing_weights[ti, k].item())
                if eid not in expert_tokens:
                    expert_tokens[eid] = ([], [])
                expert_tokens[eid][0].append(ti)
                expert_tokens[eid][1].append(w)

        # 6. 执行专家（热优先 + 冷按就绪顺序）
        output = torch.zeros_like(hs_2d)
        total_wait_s = 0.0

        def exec_expert(eid: int):
            if eid not in expert_tokens:
                return
            weights = self.store.wait_and_get(eid)
            tids, ws = expert_tokens[eid]
            out = self._run_expert(weights, hs_2d[tids])
            ws_t = torch.tensor(ws, device=self.device, dtype=hs_2d.dtype)
            output[tids] += out * ws_t.unsqueeze(-1)

        for eid in hot_set:
            exec_expert(eid)

        cold_remaining = list(cold_set)
        while cold_remaining:
            ready = [e for e in cold_remaining if self.store.is_ready(e)]
            if ready:
                for eid in ready:
                    exec_expert(eid)
                    cold_remaining.remove(eid)
            elif cold_remaining:
                eid = cold_remaining[0]
                t0 = time.perf_counter()
                exec_expert(eid)
                total_wait_s += time.perf_counter() - t0
                cold_remaining.remove(eid)

        # 7. 更新相关性表 & 状态
        if prev_r is not None:
            self.corr.update_batch(self.layer_idx, prev_r.to(self.device), routing_indices)
        self.state.update_layer(self.layer_idx, routing_indices)

        # 8. 记录指标
        self.metrics.record_layer(
            layer_idx=self.layer_idx,
            intra_bubble_s=total_wait_s,
            num_expert_used=len(active_set),
            hot_hit_count=len(hot_set),
            hot_total=max(len(hot_ids), 1),
        )

        # 9. 释放 GPU 缓存
        self.store.evict(list(active_set))

        return output.view(batch_size, seq_len, hidden_dim)

    def _run_expert(
        self, weights: Dict[str, torch.Tensor], x: torch.Tensor
    ) -> torch.Tensor:
        """
        SwiGLU FFN（与 Qwen3MoeExperts.forward 逻辑完全一致）。
          gate_up_proj : [2*intermediate, hidden]
          down_proj    : [hidden, intermediate]
        """
        dtype = x.dtype
        gate_up = weights["gate_up_proj"].to(device=self.device, dtype=dtype)
        down = weights["down_proj"].to(device=self.device, dtype=dtype)
        gate, up = F.linear(x, gate_up).chunk(2, dim=-1)
        return F.linear(self.act_fn(gate) * up, down)


# ============================================================================
#  公共接口
# ============================================================================

def build_stores(
    num_layers: int,
    num_experts: int,
    device: torch.device,
    io_stream: Optional["torch.cuda.Stream"] = None,
    compute_stream: Optional["torch.cuda.Stream"] = None,
    tracer: Optional["ExpertTracer"] = None,
) -> Dict[int, TorrentLayerExpertStore]:
    """为每个 MoE 层创建一个 TorrentLayerExpertStore。"""
    return {
        i: TorrentLayerExpertStore(
            layer_idx=i,
            num_experts=num_experts,
            device=device,
            io_stream=io_stream,
            compute_stream=compute_stream,
            tracer=tracer,
        )
        for i in range(num_layers)
    }


def _block_dequant_fp8(
    w_fp8: torch.Tensor,
    scale_inv: torch.Tensor,
    out_dtype: torch.dtype = torch.bfloat16,
) -> torch.Tensor:
    """
    块量化 FP8 反量化（per-block scale_inv，块大小 128×128）。
    """
    r, c = w_fp8.shape
    sr, sc = scale_inv.shape
    scale_tiled = (
        scale_inv.float()
        .repeat_interleave(r // sr, dim=0)
        .repeat_interleave(c // sc, dim=1)
    )
    return (w_fp8.to(torch.float32) * scale_tiled).to(out_dtype)


def _register_from_safetensors(
    model,
    stores: Dict[int, TorrentLayerExpertStore],
    model_path: str,
    dtype: torch.dtype = torch.bfloat16,
) -> int:
    """
    直接从 safetensors 文件读取 FP8 专家权重 + block scale_inv，
    正确反量化后注册。支持跨 shard 的流水线缓冲加载。
    """
    import json
    import os
    from collections import defaultdict

    try:
        from safetensors import safe_open
    except ImportError:
        raise RuntimeError("safetensors 未安装，请运行: pip install safetensors")

    idx_file = os.path.join(model_path, "model.safetensors.index.json")
    with open(idx_file) as f:
        weight_map = json.load(f)["weight_map"]

    layer_num_experts: Dict[int, int] = {}
    for layer_idx in sorted(stores.keys()):
        moe_block = model.model.layers[layer_idx].mlp
        if hasattr(moe_block, "experts"):
            layer_num_experts[layer_idx] = moe_block.experts.num_experts

    PROJECTIONS = ["gate_proj", "up_proj", "down_proj"]
    KINDS = ["weight", "weight_scale_inv"]
    needed_keys: Dict[str, Tuple[int, int, str, str]] = {}
    expert_required_keys: Dict[Tuple[int, int], set] = {}

    for layer_idx, num_exp in layer_num_experts.items():
        for eid in range(num_exp):
            pf = f"model.layers.{layer_idx}.mlp.experts.{eid}"
            keys_for_expert = set()
            for proj in PROJECTIONS:
                for kind in KINDS:
                    k = f"{pf}.{proj}.{kind}"
                    if k in weight_map:
                        needed_keys[k] = (layer_idx, eid, proj, kind)
                        keys_for_expert.add(k)
            expert_required_keys[(layer_idx, eid)] = keys_for_expert

    shard_to_keys: Dict[str, List[str]] = defaultdict(list)
    for k in needed_keys:
        shard_to_keys[weight_map[k]].append(k)

    buffer: Dict[str, torch.Tensor] = {}
    total = 0
    prev_layer = -1

    for shard_file in sorted(shard_to_keys.keys()):
        shard_path = os.path.join(model_path, shard_file)
        keys_in_shard = shard_to_keys[shard_file]
        layers_set = {needed_keys[k][0] for k in keys_in_shard}
        logger.info(
            f"  Loading shard {shard_file} "
            f"[layers {min(layers_set)}-{max(layers_set)}, "
            f"{len(keys_in_shard)} tensors] ..."
        )

        with safe_open(shard_path, framework="pt", device="cpu") as sf:
            for k in keys_in_shard:
                buffer[k] = sf.get_tensor(k)
                layer_idx, eid, proj, kind = needed_keys[k]
                expert_id = (layer_idx, eid)
                required = expert_required_keys[expert_id]

                if all(rk in buffer for rk in required):
                    pf = f"model.layers.{layer_idx}.mlp.experts.{eid}"

                    g_w = buffer[f"{pf}.gate_proj.weight"]
                    g_s = buffer[f"{pf}.gate_proj.weight_scale_inv"]
                    u_w = buffer[f"{pf}.up_proj.weight"]
                    u_s = buffer[f"{pf}.up_proj.weight_scale_inv"]
                    d_w = buffer[f"{pf}.down_proj.weight"]
                    d_s = buffer[f"{pf}.down_proj.weight_scale_inv"]

                    gate_proj = _block_dequant_fp8(g_w, g_s, dtype)
                    up_proj   = _block_dequant_fp8(u_w, u_s, dtype)
                    down_proj = _block_dequant_fp8(d_w, d_s, dtype)
                    gate_up   = torch.cat([gate_proj, up_proj], dim=0)  # [2I, H]

                    if layer_idx in stores:
                        stores[layer_idx].register(eid, gate_up, down_proj)
                        total += 1

                    for rk in required:
                        del buffer[rk]

                    if layer_idx != prev_layer:
                        if prev_layer >= 0:
                            logger.info(f"    Layer {prev_layer}: experts registered ✓")
                        prev_layer = layer_idx

    if prev_layer >= 0:
        logger.info(f"    Layer {prev_layer}: experts registered ✓")
    if buffer:
        logger.warning(f"  {len(buffer)} tensors left in buffer (incomplete experts?)")

    logger.info(f"Total experts registered from safetensors: {total}")
    return total


def register_all_experts(
    model,
    stores: Dict[int, TorrentLayerExpertStore],
    model_path: Optional[str] = None,
    dtype: torch.dtype = torch.bfloat16,
) -> int:
    """
    注册专家权重到各层 TorrentLayerExpertStore。

    - model_path 非 None：从 safetensors 直接读取（支持 FP8 量化格式）
    - model_path 为 None：直接从模型参数读（非 FP8 回退）
    """
    if model_path is not None:
        return _register_from_safetensors(model, stores, model_path, dtype)

    total = 0
    for i, layer in enumerate(model.model.layers):
        moe_block = layer.mlp
        if not hasattr(moe_block, "experts") or i not in stores:
            continue

        experts = moe_block.experts
        num_experts = experts.num_experts
        store = stores[i]

        gate_up_all = experts.gate_up_proj.data  # [E, 2I, H]
        down_all = experts.down_proj.data         # [E, H, I]

        src_dtype = gate_up_all.dtype
        needs_cast = src_dtype not in (torch.float16, torch.bfloat16, torch.float32)

        for eid in range(num_experts):
            gu = gate_up_all[eid].detach()
            dw = down_all[eid].detach()
            if needs_cast:
                gu = gu.to(torch.float32).to(dtype)
                dw = dw.to(torch.float32).to(dtype)
            elif gu.dtype != dtype:
                gu = gu.to(dtype)
                dw = dw.to(dtype)
            store.register(eid, gu.cpu(), dw.cpu())
            total += 1

        logger.info(
            f"  Layer {i}: {num_experts} experts | "
            f"gate_up={list(gate_up_all.shape[1:])} down={list(down_all.shape[1:])} | "
            f"src_dtype={src_dtype} → {dtype}"
        )

    logger.info(f"Total experts registered: {total}")
    return total


def patch_model(
    model,
    stores: Dict[int, TorrentLayerExpertStore],
    corr_table: CorrelationTable,
    state: TorrentState,
    metrics: MetricsCollector,
    device: torch.device,
    top_k: int = 8,
    layer_range: Optional[Tuple[int, int]] = None,
    tracer: Optional["ExpertTracer"] = None,
) -> Dict[int, TorrentMoEForward]:
    """
    为 Qwen3MoeForCausalLM 各层 monkey-patch Torrent MoE forward。

    Returns:
        {layer_idx: TorrentMoEForward}
    """
    layers = model.model.layers
    n_layers = len(layers)
    lo, hi = layer_range if layer_range else (0, n_layers)
    hi = min(hi, n_layers)

    hidden_act = getattr(model.config, "hidden_act", "silu")
    _ACT_MAP = {
        "silu": F.silu,
        "gelu": F.gelu,
        "relu": F.relu,
        "tanh": torch.tanh,
    }
    act_fn = _ACT_MAP.get(hidden_act, F.silu)

    torrent_forwards: Dict[int, TorrentMoEForward] = {}

    for i in range(lo, hi):
        layer = layers[i]
        moe_block = layer.mlp

        if not hasattr(moe_block, "gate") or not hasattr(moe_block, "experts"):
            continue
        if i not in stores:
            continue

        tfwd = TorrentMoEForward(
            moe_block=moe_block,
            layer_idx=i,
            store=stores[i],
            corr_table=corr_table,
            state=state,
            metrics=metrics,
            top_k=top_k,
            prefetch_k=top_k,
            device=device,
            act_fn=act_fn,
            tracer=tracer,
        )

        def make_forward(tf):
            def _forward(self_block, hs):
                return tf(hs)
            return _forward

        moe_block.forward = types.MethodType(make_forward(tfwd), moe_block)
        torrent_forwards[i] = tfwd

    logger.info(
        f"Patched {len(torrent_forwards)}/{n_layers} MoE layers with Torrent forward."
    )
    return torrent_forwards
