"""
async_loader.py
===============
Torrent 异步专家加载器（论文 §6 + §8）

职责：
  - 在 attention 计算的 CUDA stream 上并行，通过另一个 CUDA stream 异步传输权重
  - 实现"gate + 热专家先行预取 → gate 结果出来后立即加载冷专家"的两阶段预取
  - 追踪每个专家的传输状态（未请求/传输中/可用/已回收）
  - 提供 wait_expert() 阻塞接口

传输状态机：
  ABSENT → LOADING → READY → EVICTED
"""

from __future__ import annotations

import logging
import threading
import time
from enum import Enum, auto
from typing import TYPE_CHECKING, Dict, List, Optional, Set

import torch

if TYPE_CHECKING:
    from srep_moe.metrics.tracer import ExpertTracer

logger = logging.getLogger(__name__)


class ExpertState(Enum):
    ABSENT = auto()
    LOADING = auto()
    READY = auto()
    EVICTED = auto()


class AsyncLoader:
    """
    异步专家权重加载器。

    Args:
        device          : 目标 GPU 设备
        io_stream       : 专用于 I/O 的 CUDA stream
        compute_stream  : 计算使用的主 CUDA stream
        layer_idx       : 所属层索引（用于 trace 事件的 layer 字段；-1 表示未指定）
        tracer          : ExpertTracer 实例（可选；为 None 时不记录任何事件）
    """

    def __init__(
        self,
        device: torch.device,
        io_stream: Optional[torch.cuda.Stream] = None,
        compute_stream: Optional[torch.cuda.Stream] = None,
        layer_idx: int = -1,
        tracer: Optional["ExpertTracer"] = None,
    ):
        self.device = device
        self._use_cuda = (device.type == "cuda")

        if self._use_cuda:
            self.io_stream = io_stream or torch.cuda.Stream(device=device)
            self.compute_stream = compute_stream or torch.cuda.current_stream(device)
        else:
            self.io_stream = None
            self.compute_stream = None

        self.layer_idx = layer_idx
        self._tracer = tracer

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

        self._expert_state: Dict[int, ExpertState] = {}
        self._gpu_weights: Dict[int, Dict[str, torch.Tensor]] = {}
        self._transfer_events: Dict[int, torch.cuda.Event] = {}
        self._cpu_weights: Dict[int, Dict[str, torch.Tensor]] = {}
        self._transfer_start_times: Dict[int, float] = {}
        self._used_experts: Set[int] = set()

        self._lock = threading.Lock()

    # ------------------------------------------------------------------
    # Phase 设置接口
    # ------------------------------------------------------------------

    def set_phase(
        self,
        phase: str,
        step: int,
        draft_pos: Optional[int] = None,
    ) -> None:
        """
        设置当前推理阶段，由 EAGLE3SDDriver 或主推理循环调用。

        Args:
            phase     : "draft" | "verify" | "ar"
            step      : decode step 编号
            draft_pos : 草稿深度（仅 draft 阶段有意义）
        """
        self._current_phase = phase
        self._current_step = step
        self._current_draft_pos = draft_pos

    # ------------------------------------------------------------------
    # 权重注册接口
    # ------------------------------------------------------------------

    def register_expert(
        self,
        expert_id: int,
        cpu_weights: Dict[str, torch.Tensor],
    ) -> None:
        pinned = {}
        for name, w in cpu_weights.items():
            if self._use_cuda and not w.is_pinned():
                pinned[name] = w.pin_memory()
            else:
                pinned[name] = w
        with self._lock:
            self._cpu_weights[expert_id] = pinned
            self._expert_state[expert_id] = ExpertState.ABSENT

    def register_experts_batch(self, experts: Dict[int, Dict[str, torch.Tensor]]) -> None:
        for eid, weights in experts.items():
            self.register_expert(eid, weights)

    # ------------------------------------------------------------------
    # 预取接口（两阶段）
    # ------------------------------------------------------------------

    def prefetch_gate(self, gate_cpu: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]:
        if not self._use_cuda:
            return {k: v.to(self.device) for k, v in gate_cpu.items()}
        gate_gpu = {}
        with torch.cuda.stream(self.io_stream):
            for name, w in gate_cpu.items():
                gate_gpu[name] = w.to(self.device, non_blocking=True)
        return gate_gpu

    def synchronize_gate(self) -> None:
        if self._use_cuda:
            self.compute_stream.wait_stream(self.io_stream)

    def prefetch_hot_experts(self, hot_expert_ids: List[int]) -> None:
        for eid in hot_expert_ids:
            self._start_transfer(eid, context="hot")

    def prefetch_cold_experts(self, cold_expert_ids: List[int]) -> None:
        for eid in cold_expert_ids:
            self._start_transfer(eid, context="cold")

    # ------------------------------------------------------------------
    # 等待接口
    # ------------------------------------------------------------------

    def wait_expert(self, expert_id: int) -> Dict[str, torch.Tensor]:
        with self._lock:
            state = self._expert_state.get(expert_id, ExpertState.ABSENT)

        if state == ExpertState.ABSENT:
            if self._tracer:
                self._tracer.emit({
                    "event": "expert_wait_miss",
                    "expert_id": expert_id,
                    "layer": self.layer_idx,
                    "step": self._current_step,
                    "phase": self._current_phase,
                })
            logger.warning(f"Expert {expert_id} not prefetched; falling back to sync load.")
            self._start_transfer(expert_id, context="cold")
        elif state == ExpertState.READY:
            if self._tracer:
                self._tracer.emit({
                    "event": "expert_wait_hit",
                    "expert_id": expert_id,
                    "layer": self.layer_idx,
                    "step": self._current_step,
                    "phase": self._current_phase,
                })
        else:  # LOADING — need to synchronously wait
            if self._tracer:
                self._tracer.emit({
                    "event": "expert_wait_miss",
                    "expert_id": expert_id,
                    "layer": self.layer_idx,
                    "step": self._current_step,
                    "phase": self._current_phase,
                })

        if self._use_cuda:
            event = self._transfer_events.get(expert_id)
            if event is not None:
                self.compute_stream.wait_event(event)

        with self._lock:
            prev_state = self._expert_state.get(expert_id, ExpertState.ABSENT)
            self._expert_state[expert_id] = ExpertState.READY
            self._used_experts.add(expert_id)
            result = self._gpu_weights[expert_id]

        # emit prefetch_ready 当传输刚完成（LOADING → READY 的转变）
        if self._tracer and prev_state == ExpertState.LOADING:
            start_t = self._transfer_start_times.pop(expert_id, None)
            transfer_ms = (time.time() - start_t) * 1e3 if start_t is not None else 0.0
            self._tracer.emit({
                "event": "prefetch_ready",
                "expert_id": expert_id,
                "layer": self.layer_idx,
                "transfer_ms": round(transfer_ms, 4),
            })

        return result

    def is_ready(self, expert_id: int) -> bool:
        with self._lock:
            state = self._expert_state.get(expert_id, ExpertState.ABSENT)
        if state == ExpertState.READY:
            return True
        if state == ExpertState.LOADING and self._use_cuda:
            event = self._transfer_events.get(expert_id)
            if event is not None and event.query():
                with self._lock:
                    self._expert_state[expert_id] = ExpertState.READY
                return True
        return False

    # ------------------------------------------------------------------
    # 回收接口
    # ------------------------------------------------------------------

    def evict_expert(self, expert_id: int) -> None:
        with self._lock:
            was_used = expert_id in self._used_experts
            if expert_id in self._gpu_weights:
                del self._gpu_weights[expert_id]
            self._expert_state[expert_id] = ExpertState.EVICTED
            if expert_id in self._transfer_events:
                del self._transfer_events[expert_id]
            self._transfer_start_times.pop(expert_id, None)
            self._used_experts.discard(expert_id)

        if self._tracer:
            self._tracer.emit({
                "event": "expert_evict",
                "expert_id": expert_id,
                "layer": self.layer_idx,
                "was_used": was_used,
                "step": self._current_step,
                "phase": self._current_phase,
            })

    def evict_experts(self, expert_ids: List[int]) -> None:
        for eid in expert_ids:
            self.evict_expert(eid)

    def reset_layer(self) -> None:
        with self._lock:
            for eid in list(self._expert_state.keys()):
                if self._expert_state[eid] != ExpertState.LOADING:
                    self._expert_state[eid] = ExpertState.ABSENT
            self._gpu_weights.clear()
            self._transfer_events.clear()

    # ------------------------------------------------------------------
    # 内部实现
    # ------------------------------------------------------------------

    def _start_transfer(self, expert_id: int, context: str = "hot") -> None:
        with self._lock:
            state = self._expert_state.get(expert_id, ExpertState.ABSENT)
            if state in (ExpertState.LOADING, ExpertState.READY):
                return
            if expert_id not in self._cpu_weights:
                logger.error(f"Expert {expert_id} not registered; skip transfer.")
                return
            self._expert_state[expert_id] = ExpertState.LOADING
            start_t = time.time()
            self._transfer_start_times[expert_id] = start_t

        # 发出 prefetch_start 事件（在 lock 外部以避免死锁）
        if self._tracer:
            ev: Dict = {
                "event": "prefetch_start",
                "expert_id": expert_id,
                "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)

        cpu_weights = self._cpu_weights[expert_id]

        if self._use_cuda:
            event = torch.cuda.Event()
            with torch.cuda.stream(self.io_stream):
                gpu_weights = {
                    name: w.to(self.device, non_blocking=True)
                    for name, w in cpu_weights.items()
                }
                event.record(self.io_stream)
            with self._lock:
                self._gpu_weights[expert_id] = gpu_weights
                self._transfer_events[expert_id] = event
        else:
            gpu_weights = {k: v.clone() for k, v in cpu_weights.items()}
            with self._lock:
                self._gpu_weights[expert_id] = gpu_weights
                self._expert_state[expert_id] = ExpertState.READY
            # 非 CUDA 路径：传输立即完成，发出 prefetch_ready
            if self._tracer:
                start_t = self._transfer_start_times.pop(expert_id, None)
                transfer_ms = (time.time() - start_t) * 1e3 if start_t is not None else 0.0
                self._tracer.emit({
                    "event": "prefetch_ready",
                    "expert_id": expert_id,
                    "layer": self.layer_idx,
                    "transfer_ms": round(transfer_ms, 4),
                })

    def get_state_snapshot(self) -> Dict[int, str]:
        with self._lock:
            return {eid: state.name for eid, state in self._expert_state.items()}

    def get_ready_experts(self) -> List[int]:
        with self._lock:
            return [
                eid for eid, state in self._expert_state.items()
                if state == ExpertState.READY
            ]
