"""
sd_engine.py
============
EAGLE3SDDriver — 将 Eagle3DraftModel 与 TorrentNBatchEngine 对接的推测解码驱动器

Issue #2: EAGLE3SDDriver 真实草稿-验证集成

核心组件：
  HiddenStateHook   — 非侵入式提取 Qwen3 最后一层 decoder 隐状态
  EAGLE3SDDriver    — 草稿生成 + 标准 speculative sampling 验证
  SDStepRecord      — 单 step 统计记录
  SDEpisodeStats    — 多 step 汇总统计
  load_eagle3_driver — 一站式加载并返回驱动器

模型路径：
  target   : /home/sage6/.cache/huggingface/hub/
             models--Qwen--Qwen3-30B-A3B-Instruct-2507-FP8/snapshots/...
  speculator: RedHatAI/Qwen3-30B-A3B-Instruct-2507-speculator.eagle3
             （通过 SpeculatorModel.from_pretrained 加载）
"""

from __future__ import annotations

import logging
import sys
from dataclasses import dataclass, field
from pathlib import Path
from typing import TYPE_CHECKING, Any, List, Optional, Tuple

import torch

if TYPE_CHECKING:
    from srep_moe.metrics.tracer import ExpertTracer

logger = logging.getLogger(__name__)

# 模型路径常量（与 run_qwen3_moe.py 保持一致）
MODEL_PATH = (
    "/home/sage6/.cache/huggingface/hub/"
    "models--Qwen--Qwen3-30B-A3B-Instruct-2507-FP8/snapshots/"
    "5a5a776300a41aaa681dd7ff0106608ef2bc90db"
)
SPECULATOR_ID = "RedHatAI/Qwen3-30B-A3B-Instruct-2507-speculator.eagle3"


# ============================================================================
# 1. 隐状态 Hook
# ============================================================================

class HiddenStateHook:
    """
    通过 register_forward_hook 提取 Qwen3MoeForCausalLM 最后一层 decoder 的隐状态。

    不修改 model 结构——hook 在 __exit__ 时自动移除。

    用法：
        hook = HiddenStateHook(model)
        with hook:
            logits = model(input_ids).logits
        hs = hook.last_hidden_states  # shape: [1, seq_len, 2048]
    """

    def __init__(self, model: Any):
        # Qwen3 结构：model.model.layers[-1] 是最后一个 decoder layer
        self._layer = model.model.layers[-1]
        self._handle: Optional[Any] = None
        self.last_hidden_states: Optional[torch.Tensor] = None

    def __enter__(self) -> "HiddenStateHook":
        def _hook(module, input, output):
            # output[0] shape: [batch, seq_len, hidden_size]
            self.last_hidden_states = output[0].detach()

        self._handle = self._layer.register_forward_hook(_hook)
        return self

    def __exit__(self, *args: Any) -> None:
        if self._handle is not None:
            self._handle.remove()
            self._handle = None


# ============================================================================
# 2. 数据类
# ============================================================================

@dataclass
class SDStepRecord:
    """一次 speculative decoding step 的统计记录。"""
    step: int
    draft_ids: List[int]
    accepted_ids: List[int]
    tau: int    # 第一个被拒绝位置（0-indexed）；== gamma 表示全部接受
    gamma: int


@dataclass
class SDEpisodeStats:
    """SDStepRecord 列表的汇总统计。"""
    num_steps: int
    gamma: int
    rejection_positions: List[int]          # 每 step 的 tau
    realized_acceptance_rate: float         # 实测平均接受率 = mean(tau_j / gamma)
    total_draft_tokens: int
    total_accepted_tokens: int
    per_depth_acceptance: List[float]       # 长度 gamma，每深度的接受率

    @classmethod
    def from_records(
        cls,
        records: List[SDStepRecord],
        gamma: int,
    ) -> "SDEpisodeStats":
        if not records:
            return cls(
                num_steps=0, gamma=gamma,
                rejection_positions=[], realized_acceptance_rate=0.0,
                total_draft_tokens=0, total_accepted_tokens=0,
                per_depth_acceptance=[0.0] * gamma,
            )
        taus = [r.tau for r in records]
        total_draft = len(records) * gamma
        total_accepted = sum(len(r.accepted_ids) for r in records)
        per_depth = [
            sum(1 for r in records if r.tau > j) / len(records)
            for j in range(gamma)
        ]
        realized_alpha = sum(per_depth) / max(gamma, 1)
        return cls(
            num_steps=len(records),
            gamma=gamma,
            rejection_positions=taus,
            realized_acceptance_rate=realized_alpha,
            total_draft_tokens=total_draft,
            total_accepted_tokens=total_accepted,
            per_depth_acceptance=per_depth,
        )


# ============================================================================
# 3. EAGLE3SDDriver
# ============================================================================

class EAGLE3SDDriver:
    """
    将 Eagle3DraftModel 与 Torrent-patched Qwen3MoeForCausalLM 对接的 SD 驱动器。

    草稿阶段（draft）：
      1. 用 target model 做单 token decode，通过 HiddenStateHook 提取最后隐状态
      2. 将 3×hidden_states 拼接送入 Eagle3DraftModel.forward(ttt_steps=γ)
      3. 返回 γ 个草稿 token ids

    验证阶段（verify）：
      1. 将 input_ids + γ 个草稿 token 拼接，一次 target model forward
      2. 取最后 γ+1 个位置的 logits
      3. 标准 speculative sampling（Leviathan et al. 2023）

    若注入 ExpertTracer，驱动器会在 draft/verify forward 前调用
    tracer.set_phase() 以区分两阶段的 H2D 事件。

    Args:
        target_model : 已由 Torrent patch_model() 处理的 Qwen3MoeForCausalLM
        speculator   : Eagle3DraftModel 实例
        tokenizer    : Qwen3 tokenizer
        torrent_state: TorrentState（草稿与验证阶段共享）
        gamma        : 最大草稿深度（建议 γ=5）
        device       : torch.device("cuda:0")
        tracer       : ExpertTracer（可选），记录 H2D 事件并区分 phase
    """

    def __init__(
        self,
        target_model: Any,
        speculator: Any,
        tokenizer: Any,
        torrent_state: Any,
        gamma: int = 5,
        device: torch.device = torch.device("cuda:0"),
        tracer: Optional["ExpertTracer"] = None,
    ):
        self.model = target_model
        self.speculator = speculator
        self.tokenizer = tokenizer
        self.state = torrent_state
        self.gamma = gamma
        self.device = device
        self.tracer = tracer
        self._hook = HiddenStateHook(target_model)
        self._step = 0  # 全局 step 计数器

    # ------------------------------------------------------------------
    # 草稿生成
    # ------------------------------------------------------------------

    def draft(
        self,
        input_ids: torch.Tensor,        # shape: [1, seq_len]
        past_key_values: Optional[Any] = None,
    ) -> Tuple[List[int], Any]:
        """
        运行 γ 步草稿生成。

        流程：
          1. target model 做单次 forward 提取最后隐状态（H2D 记入 draft phase）
          2. 将最后 token 隐状态 × 3 拼接送入 Eagle3DraftModel（γ 步）
          3. 返回草稿 token ids 和更新的 KV cache

        Returns:
            draft_ids   : List[int]，长度 γ
            new_kv      : 更新后的 DynamicCache（仅含 context 的 KV）
        """
        if self.tracer:
            self.tracer.set_phase("draft", self._step)

        # Step 1: target model forward，提取最后层隐状态
        with self._hook:
            with torch.no_grad():
                out = self.model(
                    input_ids,
                    past_key_values=past_key_values,
                    use_cache=True,
                )
        # target_hs shape: [1, seq_len, hidden_size]
        target_hs = self._hook.last_hidden_states
        if target_hs is None:
            raise RuntimeError("HiddenStateHook 未能提取隐状态，请检查 hook 挂载。")

        # 取最后位置隐状态，拼成 3×hidden_size（EAGLE-3 fc 输入格式）
        last_hs = target_hs[:, -1:, :]              # [1, 1, H]
        eagle_input = last_hs.expand(-1, -1, 3)     # [1, 1, 3H] 逻辑上等价于 repeat
        # 严格按照 EAGLE-3 forward 的输入格式：[1, T, 3*hidden_size]
        eagle_input_hs = torch.cat([last_hs, last_hs, last_hs], dim=-1)  # [1, 1, 3H]

        # Step 2: EAGLE-3 draft forward（生成 γ 个草稿 token）
        with torch.no_grad():
            draft_result = self.speculator.forward(
                hidden_states=eagle_input_hs.to(self.speculator.device
                    if hasattr(self.speculator, "device") else self.device),
                input_ids=input_ids[:, -1:].to(self.device),
                ttt_steps=self.gamma,
            )

        # draft_result: List[Tensor[1, 1]] 长度 γ，或 (List, routing_predictions)
        if isinstance(draft_result, tuple):
            draft_tokens_tensors, _ = draft_result
        else:
            draft_tokens_tensors = draft_result

        draft_ids = [int(t[0, -1].item()) for t in draft_tokens_tensors]
        return draft_ids, out.past_key_values

    # ------------------------------------------------------------------
    # 验证（一次 target model forward）
    # ------------------------------------------------------------------

    def verify(
        self,
        input_ids: torch.Tensor,         # shape: [1, seq_len]
        draft_ids: List[int],
        past_key_values: Optional[Any] = None,
    ) -> Tuple[torch.Tensor, Any]:
        """
        将全部草稿 token 拼接到 input_ids 后，一次 target model forward 验证。

        Returns:
            verify_logits : shape [1, γ+1, vocab_size]，最后 γ+1 个位置的 logits
            new_kv        : 更新后的 DynamicCache
        """
        if self.tracer:
            self.tracer.set_phase("verify", self._step)

        draft_tensor = torch.tensor([draft_ids], dtype=torch.long, device=self.device)
        full_ids = torch.cat([input_ids, draft_tensor], dim=-1)

        with torch.no_grad():
            out = self.model(full_ids, past_key_values=past_key_values, use_cache=True)

        # 取最后 γ+1 个位置（验证 γ 个草稿 + 生成 bonus token）
        verify_logits = out.logits[:, -(self.gamma + 1):, :]
        return verify_logits, out.past_key_values

    # ------------------------------------------------------------------
    # 标准 speculative sampling（Leviathan et al. 2023）
    # ------------------------------------------------------------------

    def speculative_sampling(
        self,
        draft_ids: List[int],
        verify_logits: torch.Tensor,    # [1, γ+1, vocab_size]
        temperature: float = 1.0,
    ) -> Tuple[int, List[int]]:
        """
        判断每个草稿 token 是否被接受。

        Returns:
            tau         : 第一个被拒绝的位置（0-indexed）；== γ 表示全部接受
            accepted_ids: 最终接受序列（含 bonus / resampled token）
        """
        tau = self.gamma  # 默认：全部接受
        accepted_ids: List[int] = []

        for j, draft_id in enumerate(draft_ids):
            verify_logit = verify_logits[0, j]  # [vocab_size]

            if temperature <= 1e-6:
                # 贪婪验证：直接比较 argmax
                target_id = int(verify_logit.argmax().item())
                if target_id == draft_id:
                    accepted_ids.append(draft_id)
                else:
                    tau = j
                    accepted_ids.append(target_id)
                    break
            else:
                # 随机采样验证
                p_target = torch.softmax(verify_logit / temperature, dim=-1)
                r = torch.rand(1, device=self.device).item()
                p_accept = float(p_target[draft_id].item())
                if r < p_accept:
                    accepted_ids.append(draft_id)
                else:
                    tau = j
                    # Corrected sampling：从修正分布重采样
                    bonus = int(torch.multinomial(p_target, num_samples=1).item())
                    accepted_ids.append(bonus)
                    break
        else:
            # 全部接受：从 bonus 位置（最后一个逻辑位）采样
            bonus_logit = verify_logits[0, -1]
            if temperature <= 1e-6:
                bonus = int(bonus_logit.argmax().item())
            else:
                p_bonus = torch.softmax(bonus_logit / temperature, dim=-1)
                bonus = int(torch.multinomial(p_bonus, num_samples=1).item())
            accepted_ids.append(bonus)

        return tau, accepted_ids

    # ------------------------------------------------------------------
    # 完整 SD 推理 episode
    # ------------------------------------------------------------------

    def run_episode(
        self,
        prompts: List[str],
        max_new_tokens: int = 128,
        temperature: float = 0.0,
    ) -> SDEpisodeStats:
        """
        对给定 prompts 列表运行完整的 SD 推理，记录逐 step 统计。

        每个 prompt 独立运行，统计汇总后返回。
        若已注入 ExpertTracer，draft/verify 的 H2D 事件将携带 phase 字段。
        """
        all_stats: List[SDStepRecord] = []

        for prompt in prompts:
            input_ids = (
                self.tokenizer(prompt, return_tensors="pt")
                .input_ids.to(self.device)
            )
            past_kv = None
            step_local = 0
            initial_len = int(input_ids.shape[-1])
            self.state.reset()

            while int(input_ids.shape[-1]) < initial_len + max_new_tokens:
                self._step = step_local

                # 草稿阶段
                draft_ids, past_kv = self.draft(input_ids, past_kv)
                self.state.advance()

                # 验证阶段
                verify_logits, past_kv = self.verify(input_ids, draft_ids, past_kv)

                # 接受/拒绝判断
                tau, accepted_ids = self.speculative_sampling(
                    draft_ids, verify_logits, temperature
                )

                all_stats.append(SDStepRecord(
                    step=step_local,
                    draft_ids=draft_ids,
                    accepted_ids=accepted_ids,
                    tau=tau,
                    gamma=self.gamma,
                ))

                # 提交接受的 token，更新 input_ids
                accepted_tensor = torch.tensor(
                    [accepted_ids], dtype=torch.long, device=self.device
                )
                input_ids = torch.cat([input_ids, accepted_tensor], dim=-1)
                step_local += 1

                # EOS 检测
                eos_ids = (
                    self.tokenizer.eos_token_id
                    if isinstance(self.tokenizer.eos_token_id, list)
                    else [self.tokenizer.eos_token_id]
                )
                if any(t in eos_ids for t in accepted_ids):
                    logger.debug(f"EOS at SD step {step_local}")
                    break

                if int(input_ids.shape[-1]) >= initial_len + max_new_tokens:
                    break

        if self.tracer:
            self.tracer.set_phase("ar", self._step)  # 恢复默认 phase

        return SDEpisodeStats.from_records(all_stats, self.gamma)


# ============================================================================
# 4. 一站式加载函数
# ============================================================================

def load_eagle3_driver(
    gamma: int = 5,
    device: torch.device = torch.device("cuda:0"),
    tracer: Optional["ExpertTracer"] = None,
) -> "EAGLE3SDDriver":
    """
    一站式加载 Qwen3-30B-A3B + EAGLE-3 speculator 并返回 EAGLE3SDDriver。

    使用与 run_qwen3_moe.py 相同的模型路径和 Torrent patch 方式。

    Args:
        gamma  : 草稿深度 γ（默认 5）
        device : 目标 GPU 设备
        tracer : ExpertTracer（可选），注入后可在 experiments/srep/ 生成 trace

    Returns:
        EAGLE3SDDriver 实例
    """
    import json

    ROOT = Path(__file__).parent.parent.parent  # srep-moe/
    sys.path.insert(0, str(ROOT / "src"))

    from transformers import AutoModelForCausalLM, AutoTokenizer
    from speculators import SpeculatorModel

    from srep_moe.models import (
        TorrentState,
        build_stores,
        register_all_experts,
        patch_model,
    )
    from srep_moe.prefetch.correlation_table import CorrelationTable
    from srep_moe.metrics.collector import MetricsCollector

    # 读取模型配置
    with open(f"{MODEL_PATH}/config.json") as f:
        cfg = json.load(f)
    NUM_LAYERS = cfg["num_hidden_layers"]
    NUM_EXPERTS = cfg["num_experts"]
    TOP_K = cfg["num_experts_per_tok"]

    # CUDA streams
    io_stream: Optional[torch.cuda.Stream] = None
    compute_stream: Optional[torch.cuda.Stream] = None
    if device.type == "cuda":
        io_stream = torch.cuda.Stream(device=device)
        compute_stream = torch.cuda.current_stream(device)

    logger.info("Loading tokenizer...")
    tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH, trust_remote_code=True)

    logger.info("Loading Qwen3-30B-A3B to CPU (此过程约 60s) ...")
    target_model = AutoModelForCausalLM.from_pretrained(
        MODEL_PATH,
        torch_dtype=torch.bfloat16,
        low_cpu_mem_usage=True,
        device_map="cpu",
        trust_remote_code=True,
    )
    target_model.eval()

    # 将非专家层移到 GPU
    logger.info("Moving non-expert weights to GPU...")
    target_model.model.embed_tokens.to(device)
    target_model.model.norm.to(device)
    target_model.lm_head.to(device)
    for layer in target_model.model.layers:
        layer.self_attn.to(device)
        layer.input_layernorm.to(device)
        layer.post_attention_layernorm.to(device)
        if hasattr(layer.mlp, "gate"):
            layer.mlp.gate.to(device)

    # 构建 Torrent 专家存储并注册权重
    stores = build_stores(
        num_layers=NUM_LAYERS,
        num_experts=NUM_EXPERTS,
        device=device,
        io_stream=io_stream,
        compute_stream=compute_stream,
        tracer=tracer,
    )
    corr_table = CorrelationTable(
        num_layers=NUM_LAYERS,
        num_experts=NUM_EXPERTS,
        top_k=TOP_K,
        tracer=tracer,
    )
    torrent_state = TorrentState(num_layers=NUM_LAYERS)
    metrics = MetricsCollector(warmup_steps=2)

    logger.info("Registering expert weights from safetensors...")
    register_all_experts(target_model, stores, model_path=MODEL_PATH, dtype=torch.bfloat16)

    # Patch 模型
    logger.info("Patching MoE layers with Torrent forward...")
    patch_model(
        model=target_model,
        stores=stores,
        corr_table=corr_table,
        state=torrent_state,
        metrics=metrics,
        device=device,
        top_k=TOP_K,
        tracer=tracer,
    )

    # 加载 EAGLE-3 speculator
    logger.info("Loading EAGLE-3 speculator...")
    speculator = SpeculatorModel.from_pretrained(
        SPECULATOR_ID,
        verifier=target_model,
        torch_dtype=torch.bfloat16,
    )
    speculator = speculator.to(device).eval()

    return EAGLE3SDDriver(
        target_model=target_model,
        speculator=speculator,
        tokenizer=tokenizer,
        torrent_state=torrent_state,
        gamma=gamma,
        device=device,
        tracer=tracer,
    )
