"""
Memory Manager - GPU and CPU memory tracking and management.

Provides utilities for:
- Querying available GPU/CPU memory
- Estimating memory requirements for model components
- Automatic memory budget calculation
- GPU memory cleanup
- CPU RAM feasibility checks for mmap mode
"""

from __future__ import annotations

import gc
import logging
from dataclasses import dataclass
from typing import Optional

import torch

logger = logging.getLogger(__name__)


@dataclass
class MemoryState:
    """Snapshot of current memory usage."""

    gpu_total_bytes: int = 0
    gpu_used_bytes: int = 0
    gpu_free_bytes: int = 0
    gpu_reserved_bytes: int = 0
    cpu_total_bytes: int = 0
    cpu_available_bytes: int = 0
    cpu_used_bytes: int = 0

    @property
    def gpu_utilization(self) -> float:
        if self.gpu_total_bytes == 0:
            return 0.0
        return self.gpu_used_bytes / self.gpu_total_bytes

    def summary(self) -> str:
        lines = []
        if self.gpu_total_bytes > 0:
            lines.append(
                f"GPU: {self.gpu_used_bytes / 1e9:.2f} / "
                f"{self.gpu_total_bytes / 1e9:.2f} GB "
                f"({self.gpu_utilization:.0%} used), "
                f"{self.gpu_free_bytes / 1e9:.2f} GB free"
            )
        lines.append(
            f"CPU: {self.cpu_used_bytes / 1e9:.2f} / "
            f"{self.cpu_total_bytes / 1e9:.2f} GB, "
            f"{self.cpu_available_bytes / 1e9:.2f} GB available"
        )
        return "\n".join(lines)


class MemoryManager:
    """
    Manages memory budgets for the inference engine.

    Handles both GPU and CPU memory, providing automatic budget
    calculation and cleanup utilities.

    CPU mmap mode note:
      In mmap mode the OS manages weight RAM via page cache. The
      "budget" concept changes: instead of Python-managed tensor RAM,
      we care that the ACTIVE computation fit in RAM (one layer at a
      time + activations + KV cache). The rest is handled by the OS.
    """

    # Reserve fraction of GPU memory for PyTorch + KV cache + activations
    GPU_RESERVE_FRACTION = 0.15
    GPU_RESERVE_MIN_BYTES = 512 * 1024**2  # 512 MB

    # Reserve for CPU mode: OS + Python overhead + KV cache + activations
    # In mmap mode we only need enough RAM for ONE layer's worth of activations
    # at a time, so reserve generously for safety.
    CPU_RESERVE_BYTES = 1 * 1024**3  # 1 GB

    def __init__(self, device: Optional[torch.device] = None):
        self.device = device or self._detect_device()
        self._initial_state = self.get_state()

    @staticmethod
    def _detect_device() -> torch.device:
        """Auto-detect the best available device."""
        if torch.cuda.is_available():
            return torch.device("cuda:0")
        if hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
            return torch.device("mps")
        return torch.device("cpu")

    def get_state(self) -> MemoryState:
        """Get current memory state."""
        state = MemoryState()

        try:
            import psutil
            vm = psutil.virtual_memory()
            state.cpu_total_bytes = vm.total
            state.cpu_available_bytes = vm.available
            state.cpu_used_bytes = vm.used
        except ImportError:
            # Fallback: try /proc/meminfo
            try:
                with open("/proc/meminfo") as f:
                    info = {}
                    for line in f:
                        parts = line.split()
                        if len(parts) >= 2:
                            info[parts[0].rstrip(":")] = int(parts[1]) * 1024
                state.cpu_total_bytes = info.get("MemTotal", 0)
                state.cpu_available_bytes = info.get("MemAvailable", 0)
                state.cpu_used_bytes = (
                    state.cpu_total_bytes - state.cpu_available_bytes
                )
            except OSError:
                pass

        if self.device.type == "cuda":
            try:
                idx = self.device.index or 0
                props = torch.cuda.get_device_properties(idx)
                state.gpu_total_bytes = props.total_memory
                state.gpu_reserved_bytes = torch.cuda.memory_reserved(idx)
                state.gpu_used_bytes = torch.cuda.memory_allocated(idx)
                state.gpu_free_bytes = (
                    state.gpu_total_bytes - state.gpu_used_bytes
                )
            except Exception as e:
                logger.warning("Could not query GPU memory: %s", e)

        return state

    def get_weight_budget_bytes(self, use_mmap: bool = False) -> int:
        """
        Calculate memory budget for cached model weights.

        GPU mode: bytes of VRAM available for the layer cache.
        CPU mmap mode: returns 0 — no Python-level layer cache is needed
          since the OS page cache handles weight caching automatically.
        CPU non-mmap mode: bytes of RAM usable as an explicit weight cache.
        """
        if use_mmap:
            # In mmap mode, no Python-managed weight cache is needed.
            # The OS page cache acts as a free, smart weight cache.
            return 0

        state = self.get_state()

        if self.device.type == "cpu":
            # Leave at least CPU_RESERVE_BYTES for OS + activations + KV cache
            available = max(
                state.cpu_available_bytes - self.CPU_RESERVE_BYTES, 0
            )
            # Use at most 60% of available RAM to avoid swapping
            budget = int(available * 0.6)
            logger.info(
                "CPU weight cache budget: %.2f GB "
                "(available: %.2f GB, reserve: %.2f GB)",
                budget / 1e9,
                state.cpu_available_bytes / 1e9,
                self.CPU_RESERVE_BYTES / 1e9,
            )
            return budget

        # GPU mode
        reserve = max(
            int(state.gpu_total_bytes * self.GPU_RESERVE_FRACTION),
            self.GPU_RESERVE_MIN_BYTES,
        )
        budget = max(state.gpu_free_bytes - reserve, 0)
        logger.info(
            "GPU weight budget: %.2f GB (free: %.2f GB, reserve: %.2f GB)",
            budget / 1e9,
            state.gpu_free_bytes / 1e9,
            reserve / 1e9,
        )
        return budget

    def get_cpu_buffer_budget_bytes(self) -> int:
        """
        CPU RAM budget for GPU-mode prefetch buffer.

        Not used in CPU mmap mode.

        Conservative sizing: large language models (especially MoE) stream
        many expert weight tensors through the CPU buffer.  On machines with
        less than ~64 GB RAM we must be careful not to crowd out the OS page
        cache and other processes.  Using more than 25 % of available RAM here
        risks triggering the OOM-killer on low-memory systems.
        """
        state = self.get_state()
        available = state.cpu_available_bytes
        # 25 % of available RAM, capped at 8 GB, at least 512 MB.
        budget = min(available // 4, 8 * 1024**3)
        budget = max(budget, 512 * 1024**2)
        logger.info(
            "CPU prefetch buffer budget: %.2f GB "
            "(available RAM: %.2f GB)",
            budget / 1e9,
            available / 1e9,
        )
        return budget

    def estimate_kv_cache_bytes(
        self,
        num_layers: int,
        hidden_size: int,
        num_heads: int,
        max_seq_len: int,
        dtype_bytes: int = 2,
    ) -> int:
        """Estimate KV cache memory (stays in active memory always)."""
        head_dim = hidden_size // max(num_heads, 1)
        per_token = 2 * num_layers * num_heads * head_dim * dtype_bytes
        return per_token * max_seq_len

    def estimate_activation_bytes(
        self,
        hidden_size: int,
        batch_size: int = 1,
        seq_len: int = 1,
        dtype_bytes: int = 4,  # float32 for CPU
    ) -> int:
        """Estimate peak activation memory during a single layer's forward."""
        per_token = 4 * hidden_size * dtype_bytes
        return per_token * batch_size * seq_len

    def cleanup_gpu(self) -> int:
        """Free unused GPU memory. Returns bytes freed."""
        if self.device.type != "cuda":
            return 0
        before = torch.cuda.memory_allocated(self.device)
        gc.collect()
        torch.cuda.empty_cache()
        after = torch.cuda.memory_allocated(self.device)
        freed = max(before - after, 0)
        if freed > 0:
            logger.info("Freed %.1f MB GPU memory", freed / 1e6)
        return freed

    def check_feasibility(
        self,
        max_layer_size_bytes: int,
        embedding_size_bytes: int,
        head_size_bytes: int,
        use_mmap: bool = False,
    ) -> tuple[bool, str]:
        """
        Check if the model can run on available hardware.

        GPU mode: needs enough VRAM for one layer + embedding + head.
        CPU mmap mode: needs enough RAM for one layer's activations only
          (weights come from mmap — OS manages paging).
        CPU non-mmap mode: needs enough RAM for one layer + embedding + head.
        """
        state = self.get_state()

        if self.device.type == "cpu":
            if use_mmap:
                # mmap mode: only active computation needs RAM.
                # The largest RAM need is during the embedding lookup or
                # single-layer attention. Estimate conservatively.
                # Activations (fp32, seq=512, hidden=4096): ~8 MB
                # KV cache (70B, seq=512): ~60 MB
                # Overhead: 1 GB
                # Total: ~1.2 GB active at any time
                active_needed = self.CPU_RESERVE_BYTES  # ~1 GB active overhead
                available = state.cpu_available_bytes

                if available < active_needed:
                    return False, (
                        f"CPU mmap mode needs ~{active_needed / 1e9:.1f} GB "
                        f"active RAM (for activations + KV cache + OS), "
                        f"but only {available / 1e9:.1f} GB available. "
                        f"Close other applications and try again."
                    )

                return True, (
                    f"CPU mmap mode: FEASIBLE with any model size. "
                    f"Available RAM: {available / 1e9:.1f} GB. "
                    f"Weights served from OS page cache (full FP16 accuracy). "
                    f"Speed: {self._estimate_cpu_tok_per_sec(max_layer_size_bytes)}"
                )
            else:
                # Standard CPU mode: need one layer + embedding + head in RAM
                required = (
                    max_layer_size_bytes
                    + embedding_size_bytes
                    + head_size_bytes
                    + self.CPU_RESERVE_BYTES
                )
                available = state.cpu_available_bytes

                if required > available:
                    return False, (
                        f"CPU mode needs ~{required / 1e9:.2f} GB RAM "
                        f"(layer: {max_layer_size_bytes / 1e9:.2f} GB + "
                        f"emb: {embedding_size_bytes / 1e9:.2f} GB + "
                        f"head: {head_size_bytes / 1e9:.2f} GB + reserve), "
                        f"but only {available / 1e9:.2f} GB available. "
                        f"Use use_mmap=True to run on less RAM."
                    )

                return True, (
                    f"CPU mode: feasible. "
                    f"Required: {required / 1e9:.2f} GB, "
                    f"Available: {available / 1e9:.2f} GB"
                )

        # GPU mode
        required = (
            max_layer_size_bytes
            + embedding_size_bytes
            + head_size_bytes
            + self.GPU_RESERVE_MIN_BYTES
        )
        if required > state.gpu_total_bytes:
            return False, (
                f"Model requires {required / 1e9:.2f} GB VRAM "
                f"(layer: {max_layer_size_bytes / 1e9:.2f} GB, "
                f"emb: {embedding_size_bytes / 1e9:.2f} GB), "
                f"but GPU only has {state.gpu_total_bytes / 1e9:.2f} GB. "
                f"Try compression='4bit' to reduce layer sizes."
            )

        return True, (
            f"GPU mode: feasible. "
            f"Required: {required / 1e9:.2f} GB, "
            f"Available: {state.gpu_total_bytes / 1e9:.2f} GB"
        )

    @staticmethod
    def _estimate_cpu_tok_per_sec(max_layer_size_bytes: int) -> str:
        """
        Rough speed estimate for CPU mmap mode based on layer size.

        Assumes a fast NVMe SSD (~3 GB/s read) as the bottleneck.
        """
        ssd_bw = 3e9  # 3 GB/s (typical NVMe)
        # I/O per token ≈ all layers (rough; cache helps for hot layers)
        sec_per_token_io_bound = max_layer_size_bytes / ssd_bw
        if sec_per_token_io_bound < 1:
            return f"~{1/sec_per_token_io_bound:.1f} tok/s (SSD-bound estimate)"
        else:
            return f"~{sec_per_token_io_bound:.0f}s/token (SSD-bound estimate)"
