"""
memory_manager.py
=================
Torrent 三层异构存储自适应放置管理器（论文 §6.1）

存储层次：
  GPU VRAM  → 当前计算 + 预取数据 + 热点常驻张量
  CPU DRAM  → 专家权重主存储（低延迟响应 gate 触发的请求）
  Disk      → 超出 DRAM 容量时的溢出存储
"""

from __future__ import annotations

import logging
import time
from collections import OrderedDict
from dataclasses import dataclass, field
from enum import Enum, auto
from typing import Dict, List, Optional, Set

import torch
import psutil

logger = logging.getLogger(__name__)


class StorageTier(Enum):
    GPU = "gpu"
    CPU = "cpu"
    DISK = "disk"
    UNKNOWN = "unknown"


@dataclass
class TensorRecord:
    """追踪单个权重张量的存储状态。"""
    key: str
    size_bytes: int
    tier: StorageTier
    last_access_time: float = 0.0
    access_count: int = 0
    is_pinned: bool = False
    is_resident: bool = False


class MemoryManager:
    """
    三层异构存储管理器。

    Args:
        gpu_budget_gb   : 分配给 Torrent 的 GPU 显存预算（GB）
        cpu_budget_gb   : 分配给 Torrent 的 CPU 内存预算（GB）
        device          : 目标 GPU 设备
        safety_margin   : 每层的安全余量比例
    """

    def __init__(
        self,
        gpu_budget_gb: float = 4.0,
        cpu_budget_gb: float = 32.0,
        device: Optional[torch.device] = None,
        safety_margin: float = 0.05,
    ):
        self.gpu_budget = int(gpu_budget_gb * 1e9)
        self.cpu_budget = int(cpu_budget_gb * 1e9)
        self.device = device or torch.device("cuda" if torch.cuda.is_available() else "cpu")
        self.safety_margin = safety_margin

        self._gpu_used: int = 0
        self._cpu_used: int = 0

        self._records: OrderedDict[str, TensorRecord] = OrderedDict()
        self._gpu_keys: Set[str] = set()
        self._cpu_keys: Set[str] = set()

        self._io_bytes_cpu_to_gpu: int = 0
        self._io_bytes_gpu_to_cpu: int = 0
        self._peak_gpu_bytes: int = 0

    # ------------------------------------------------------------------
    # 核心接口
    # ------------------------------------------------------------------

    def can_fit_on_gpu(self, size_bytes: int) -> bool:
        effective = int(self.gpu_budget * (1 - self.safety_margin))
        return self._gpu_used + size_bytes <= effective

    def can_fit_on_cpu(self, size_bytes: int) -> bool:
        effective = int(self.cpu_budget * (1 - self.safety_margin))
        return self._cpu_used + size_bytes <= effective

    def register_tensor(
        self,
        key: str,
        tensor: torch.Tensor,
        tier: StorageTier = StorageTier.CPU,
        is_resident: bool = False,
    ) -> TensorRecord:
        size = tensor.nelement() * tensor.element_size()
        rec = TensorRecord(
            key=key,
            size_bytes=size,
            tier=tier,
            last_access_time=time.monotonic(),
            is_pinned=tensor.is_pinned() if self.device.type == "cuda" else False,
            is_resident=is_resident,
        )
        self._records[key] = rec

        if tier == StorageTier.GPU:
            self._gpu_used += size
            self._gpu_keys.add(key)
            self._peak_gpu_bytes = max(self._peak_gpu_bytes, self._gpu_used)
        elif tier == StorageTier.CPU:
            self._cpu_used += size
            self._cpu_keys.add(key)

        return rec

    def mark_loaded_to_gpu(self, key: str) -> None:
        rec = self._records.get(key)
        if rec is None:
            return
        if rec.tier == StorageTier.CPU:
            self._cpu_used -= rec.size_bytes
            self._cpu_keys.discard(key)
        self._io_bytes_cpu_to_gpu += rec.size_bytes
        self._gpu_used += rec.size_bytes
        self._gpu_keys.add(key)
        rec.tier = StorageTier.GPU
        rec.last_access_time = time.monotonic()
        rec.access_count += 1
        self._peak_gpu_bytes = max(self._peak_gpu_bytes, self._gpu_used)

    def mark_evicted_from_gpu(self, key: str) -> None:
        rec = self._records.get(key)
        if rec is None or rec.tier != StorageTier.GPU:
            return
        self._gpu_used -= rec.size_bytes
        self._gpu_keys.discard(key)
        rec.tier = StorageTier.CPU
        self._cpu_used += rec.size_bytes
        self._cpu_keys.add(key)
        self._io_bytes_gpu_to_cpu += rec.size_bytes

    # ------------------------------------------------------------------
    # LRU 淘汰策略
    # ------------------------------------------------------------------

    def evict_lru_from_gpu(self, required_bytes: int) -> List[str]:
        evicted = []
        candidates = [
            (rec.last_access_time, key)
            for key, rec in self._records.items()
            if rec.tier == StorageTier.GPU and not rec.is_resident
        ]
        candidates.sort()

        freed = 0
        for _, key in candidates:
            if freed >= required_bytes:
                break
            self.mark_evicted_from_gpu(key)
            evicted.append(key)
            freed += self._records[key].size_bytes

        return evicted

    def place_tensor_adaptive(
        self,
        key: str,
        tensor: torch.Tensor,
        prefer_gpu: bool = False,
    ) -> StorageTier:
        size = tensor.nelement() * tensor.element_size()
        if prefer_gpu and self.can_fit_on_gpu(size):
            self.register_tensor(key, tensor, StorageTier.GPU, is_resident=prefer_gpu)
            return StorageTier.GPU
        elif self.can_fit_on_cpu(size):
            self.register_tensor(key, tensor, StorageTier.CPU)
            return StorageTier.CPU
        else:
            logger.warning(f"Tensor {key} ({size/1e9:.2f}GB) cannot fit in CPU; would need disk.")
            self.register_tensor(key, tensor, StorageTier.DISK)
            return StorageTier.DISK

    # ------------------------------------------------------------------
    # 统计接口
    # ------------------------------------------------------------------

    def get_stats(self) -> dict:
        return {
            "gpu_used_gb": self._gpu_used / 1e9,
            "gpu_budget_gb": self.gpu_budget / 1e9,
            "gpu_utilization": self._gpu_used / max(self.gpu_budget, 1),
            "cpu_used_gb": self._cpu_used / 1e9,
            "cpu_budget_gb": self.cpu_budget / 1e9,
            "peak_gpu_gb": self._peak_gpu_bytes / 1e9,
            "io_cpu_to_gpu_gb": self._io_bytes_cpu_to_gpu / 1e9,
            "io_gpu_to_cpu_gb": self._io_bytes_gpu_to_cpu / 1e9,
            "total_tensors": len(self._records),
            "gpu_tensors": len(self._gpu_keys),
            "cpu_tensors": len(self._cpu_keys),
        }

    def reset_io_stats(self) -> None:
        self._io_bytes_cpu_to_gpu = 0
        self._io_bytes_gpu_to_cpu = 0
        self._peak_gpu_bytes = 0
