"""
Tensor Streaming Engine - Async I/O with sub-layer pipelining.

Supports two loading modes:

  GPU MODE (default when CUDA/MPS available):
    ┌──────────┐    ┌──────────┐    ┌──────────┐
    │  Disk    │───▶│  CPU     │───▶│  GPU     │
    │  (SSD)   │    │  Buffer  │    │  Compute │
    └──────────┘    └──────────┘    └──────────┘
    3-level async pipeline: prefetch → transfer → compute

  CPU MMAP MODE (default when CPU-only):
    ┌──────────────────────────────────────────┐
    │  Disk (memory-mapped by OS)              │
    │  Pages loaded on-demand, cached by OS   │
    └──────────────────────────────────────────┘
    Zero-copy: tensor data lives in the OS page cache, not Python heap.
    - Full FP16 accuracy (no quantization needed)
    - OS page cache = automatic smart caching of hot layers
    - Works with as little as 4 GB RAM for any model size
    - The OS evicts pages when RAM is needed, loads them back on access

The pipeline operates at three levels (GPU mode):
1. Layer-level: While layer N computes, layer N+1 loads to CPU
2. Component-level: While attention computes, MLP loads
3. Tensor-level: Weight matrices stream in blocks (future optimization)
"""

from __future__ import annotations

import logging
import threading
from collections import OrderedDict
from concurrent.futures import Future, ThreadPoolExecutor
from pathlib import Path
from typing import Optional

import torch

from dokodemo_ai.graph.compiler import ComponentInfo, LayerInfo, ModelGraph, TensorRef

logger = logging.getLogger(__name__)


class TensorStreamer:
    """
    Asynchronous tensor loading with multi-level prefetching (GPU mode)
    or zero-copy memory-mapped loading (CPU mode).

    In CPU mmap mode:
    - Tensors are backed by the OS page cache (zero Python-side copy)
    - The OS transparently loads/evicts pages based on usage
    - No explicit CPU buffer needed — the OS IS the buffer
    - Any model size works on 4 GB RAM with full FP16 accuracy
    """

    def __init__(
        self,
        model_graph: ModelGraph,
        device: torch.device,
        num_workers: int = 2,
        cpu_buffer_bytes: int = 0,
        pin_memory: bool = True,
        use_mmap: bool = False,
    ):
        """
        Args:
            model_graph: Compiled model graph with SafeTensors references.
            device: Target device for computation.
            num_workers: Number of I/O threads for prefetching (GPU mode).
            cpu_buffer_bytes: Max CPU memory for prefetch buffer (GPU mode only).
            pin_memory: Use pinned memory for faster CPU→GPU transfers.
            use_mmap: Use memory-mapped loading (recommended for CPU-only).
                     When True, tensors are backed by OS page cache — no
                     explicit CPU RAM allocation needed.
        """
        self.graph = model_graph
        self.device = device
        self.use_mmap = use_mmap
        self.pin_memory = pin_memory and device.type == "cuda"

        # Thread pool for async I/O (GPU mode)
        self._executor = ThreadPoolExecutor(
            max_workers=num_workers if not use_mmap else 1,
            thread_name_prefix="dokodemo-io",
        )

        if use_mmap:
            # CPU mmap mode: no Python-level buffer needed
            self._cpu_buffer: OrderedDict[str, dict[str, torch.Tensor]] = (
                OrderedDict()
            )
            self._cpu_buffer_bytes = 0
            self._current_cpu_bytes = 0
            logger.info(
                "TensorStreamer (mmap mode): OS page cache manages weight RAM. "
                "Full FP16 accuracy, works on 4GB RAM."
            )
        else:
            # GPU mode: maintain explicit CPU prefetch buffer
            self._cpu_buffer = OrderedDict()
            self._cpu_buffer_bytes = cpu_buffer_bytes or self._auto_cpu_budget()
            self._current_cpu_bytes = 0
            logger.info(
                "TensorStreamer (GPU mode): %d workers, %.1f GB CPU buffer",
                num_workers,
                self._cpu_buffer_bytes / 1e9,
            )

        # Track in-flight prefetch futures
        self._prefetch_futures: dict[str, Future] = {}
        self._lock = threading.Lock()

        # SafeTensors file handles
        self._file_handles: dict[str, object] = {}
        self._file_handles_lock = threading.Lock()

    def _auto_cpu_budget(self) -> int:
        """Estimate a conservative CPU buffer size (GPU mode only).

        Consistent with MemoryManager.get_cpu_buffer_budget_bytes():
        25% of available RAM, capped at 8 GB, at least 512 MB.
        """
        try:
            import psutil
            available = psutil.virtual_memory().available
            budget = min(available // 4, 8 * 1024**3)
            return max(budget, 512 * 1024**2)
        except ImportError:
            return 2 * 1024**3

    # ------------------------------------------------------------------
    # Synchronous loading (blocking)
    # ------------------------------------------------------------------

    def load_component(
        self,
        component: ComponentInfo,
        to_device: Optional[torch.device] = None,
    ) -> dict[str, torch.Tensor]:
        """
        Load a component's weights, using cache/mmap if available.

        In mmap mode: returns OS-page-backed tensors (zero RAM allocation).
        In GPU mode: checks CPU buffer, then loads from disk if needed.
        """
        target = to_device or self.device

        if self.use_mmap:
            return self._load_mmap(component.params, target)

        # GPU mode: check CPU buffer first
        cache_key = component.name
        with self._lock:
            if cache_key in self._cpu_buffer:
                self._cpu_buffer.move_to_end(cache_key)
                cpu_tensors = self._cpu_buffer[cache_key]
                return self._transfer_to_device(cpu_tensors, target)

        # Check in-flight prefetch
        future = self._prefetch_futures.get(cache_key)
        if future is not None:
            cpu_tensors = future.result()
            with self._lock:
                self._prefetch_futures.pop(cache_key, None)
            return self._transfer_to_device(cpu_tensors, target)

        # Synchronous load from disk
        cpu_tensors = self._load_tensors_to_cpu(component.params)
        return self._transfer_to_device(cpu_tensors, target)

    def load_layer(
        self,
        layer: LayerInfo,
        to_device: Optional[torch.device] = None,
        exclude_experts: bool = False,
    ) -> dict[str, torch.Tensor]:
        """Load all weights for a layer."""
        all_tensors: dict[str, torch.Tensor] = {}
        target = to_device or self.device

        for comp in layer.components:
            if exclude_experts and comp.kind == "expert":
                continue
            comp_tensors = self.load_component(comp, to_device=target)
            all_tensors.update(comp_tensors)

        return all_tensors

    def load_expert(
        self,
        layer: LayerInfo,
        expert_index: int,
        to_device: Optional[torch.device] = None,
    ) -> dict[str, torch.Tensor]:
        """Load a single expert's weights."""
        experts = layer.experts
        if expert_index >= len(experts):
            raise IndexError(
                f"Expert index {expert_index} out of range "
                f"(layer has {len(experts)} experts)"
            )
        return self.load_component(experts[expert_index], to_device)

    # ------------------------------------------------------------------
    # Asynchronous prefetching (GPU mode only)
    # ------------------------------------------------------------------

    def prefetch_component(self, component: ComponentInfo) -> None:
        """Start loading a component to CPU buffer in the background."""
        if self.use_mmap:
            return  # mmap: OS handles prefetching via read-ahead

        cache_key = component.name
        with self._lock:
            if cache_key in self._cpu_buffer:
                return
            if cache_key in self._prefetch_futures:
                return

        future = self._executor.submit(self._prefetch_to_cpu, component)
        with self._lock:
            self._prefetch_futures[cache_key] = future

    def prefetch_layer(
        self, layer: LayerInfo, exclude_experts: bool = False
    ) -> None:
        """Start loading a layer's components to CPU buffer in background."""
        if self.use_mmap:
            return  # OS handles it
        for comp in layer.components:
            if exclude_experts and comp.kind == "expert":
                continue
            self.prefetch_component(comp)

    def prefetch_experts(
        self, layer: LayerInfo, expert_indices: list[int]
    ) -> None:
        """Prefetch specific experts for a MoE layer."""
        if self.use_mmap:
            return  # OS handles it
        experts = layer.experts
        for idx in expert_indices:
            if idx < len(experts):
                self.prefetch_component(experts[idx])

    # ------------------------------------------------------------------
    # mmap loading (CPU mode) — zero-copy, OS-managed RAM
    # ------------------------------------------------------------------

    def _load_mmap(
        self,
        params: list[TensorRef],
        target: torch.device,
    ) -> dict[str, torch.Tensor]:
        """
        Load tensors using OS memory mapping — zero Python-side copy.

        How it works:
        1. SafeTensors files are memory-mapped by the OS
        2. We get a numpy view of the mmap region (no copy)
        3. torch.from_numpy() shares the memory (no copy)
        4. The tensor's data lives in the OS page cache
        5. When RAM is full, the OS evicts pages to disk automatically
        6. When we access evicted data, the OS loads it back

        This means we can run a 405B model (810 GB FP16) on 4 GB RAM
        with full accuracy — the OS manages which weights are resident.

        Computation note: weights are stored as FP16 in SafeTensors.
        PyTorch automatically upscales to float32 during matmul on CPUs
        that don't support native FP16, so accuracy is preserved.
        """
        import numpy as np

        tensors: dict[str, torch.Tensor] = {}

        for ref in params:
            handle = self._get_mmap_handle(ref.file)
            # Get numpy array backed by OS mmap (zero-copy)
            np_arr = handle.get_tensor(ref.name)
            # torch.from_numpy shares memory — still zero-copy
            tensor = torch.from_numpy(np_arr)

            if target.type != "cpu":
                # For GPU targets even in mmap mode, we still need to transfer
                tensor = tensor.to(target, non_blocking=True)
            # else: tensor stays as mmap-backed CPU tensor

            tensors[ref.name] = tensor

        return tensors

    def _get_mmap_handle(self, filename: str):
        """Get or open a SafeTensors file handle for mmap/numpy access."""
        with self._file_handles_lock:
            if filename not in self._file_handles:
                from safetensors import safe_open
                filepath = Path(self.graph.model_path) / filename
                # framework="numpy" gives mmap-backed numpy arrays
                self._file_handles[filename] = safe_open(
                    str(filepath), framework="numpy"
                )
            return self._file_handles[filename]

    # ------------------------------------------------------------------
    # Standard (non-mmap) I/O operations (GPU mode)
    # ------------------------------------------------------------------

    def _get_file_handle(self, filename: str):
        """Get or create a SafeTensors file handle (PyTorch mode)."""
        with self._file_handles_lock:
            if filename not in self._file_handles:
                from safetensors import safe_open
                filepath = Path(self.graph.model_path) / filename
                self._file_handles[filename] = safe_open(
                    str(filepath), framework="pt"
                )
            return self._file_handles[filename]

    # Tensors larger than this threshold are NOT pinned — pinning doubles the
    # temporary RAM usage (original + locked copy), which can OOM on 16 GB
    # machines when loading 350 MB+ Mixtral expert tensors.
    _PIN_MEMORY_LIMIT_BYTES = 100 * 1024 * 1024  # 100 MB

    def _load_tensors_to_cpu(
        self, params: list[TensorRef]
    ) -> dict[str, torch.Tensor]:
        """Load tensors from disk to CPU memory (GPU mode)."""
        tensors: dict[str, torch.Tensor] = {}

        for ref in params:
            handle = self._get_file_handle(ref.file)
            tensor = handle.get_tensor(ref.name)

            if (
                self.pin_memory
                and tensor.device.type == "cpu"
                and tensor.nbytes < self._PIN_MEMORY_LIMIT_BYTES
            ):
                tensor = tensor.pin_memory()

            tensors[ref.name] = tensor

        return tensors

    def _prefetch_to_cpu(self, component: ComponentInfo) -> dict[str, torch.Tensor]:
        """Background task: load component to CPU buffer."""
        tensors = self._load_tensors_to_cpu(component.params)

        with self._lock:
            comp_size = sum(t.nbytes for t in tensors.values())
            self._evict_cpu_buffer(comp_size)
            self._cpu_buffer[component.name] = tensors
            self._current_cpu_bytes += comp_size

        return tensors

    def _transfer_to_device(
        self,
        cpu_tensors: dict[str, torch.Tensor],
        device: torch.device,
    ) -> dict[str, torch.Tensor]:
        """Transfer tensors from CPU to target device."""
        if device.type == "cpu":
            return cpu_tensors

        gpu_tensors: dict[str, torch.Tensor] = {}
        for name, tensor in cpu_tensors.items():
            gpu_tensors[name] = tensor.to(device, non_blocking=True)
        return gpu_tensors

    def _evict_cpu_buffer(self, needed_bytes: int) -> None:
        """Evict old entries from CPU buffer to make room (holds _lock)."""
        while (
            self._current_cpu_bytes + needed_bytes > self._cpu_buffer_bytes
            and self._cpu_buffer
        ):
            key, tensors = self._cpu_buffer.popitem(last=False)
            freed = sum(t.nbytes for t in tensors.values())
            self._current_cpu_bytes -= freed
            del tensors

    # ------------------------------------------------------------------
    # Direct name-based loading (used by VLM vision encoder init)
    # ------------------------------------------------------------------

    def load_tensors_by_name(
        self,
        param_names: list[str],
        to_device: Optional[torch.device] = None,
    ) -> dict[str, torch.Tensor]:
        """Load tensors by parameter name, using safetensors_index for file lookup.

        Used to load vision encoder + projector weights at startup for VLMs.
        These tensors stay resident in GPU memory for the model's lifetime.
        """
        target = to_device or self.device
        index = self.graph.safetensors_index

        # Build minimal TensorRef objects from the index
        refs = []
        for name in param_names:
            if name in index:
                refs.append(TensorRef(name=name, file=index[name], dtype="F16"))
            else:
                logger.warning("Vision param '%s' not in safetensors index, skipping", name)

        if not refs:
            return {}

        if self.use_mmap:
            return self._load_mmap(refs, target)

        cpu_tensors = self._load_tensors_to_cpu(refs)
        return self._transfer_to_device(cpu_tensors, target)

    # ------------------------------------------------------------------
    # Cleanup
    # ------------------------------------------------------------------

    def evict_component(self, component_name: str) -> None:
        """Manually evict a component from CPU buffer."""
        if self.use_mmap:
            return  # Nothing to evict — OS manages it
        with self._lock:
            if component_name in self._cpu_buffer:
                tensors = self._cpu_buffer.pop(component_name)
                self._current_cpu_bytes -= sum(t.nbytes for t in tensors.values())
                del tensors

    def evict_layer(self, layer: LayerInfo) -> None:
        """Evict all components of a layer from CPU buffer."""
        for comp in layer.components:
            self.evict_component(comp.name)

    def shutdown(self) -> None:
        """Shutdown the streamer, closing all handles and threads."""
        self._executor.shutdown(wait=False)

        with self._file_handles_lock:
            for handle in self._file_handles.values():
                if hasattr(handle, "__exit__"):
                    try:
                        handle.__exit__(None, None, None)
                    except Exception:
                        pass
            self._file_handles.clear()

        with self._lock:
            self._cpu_buffer.clear()
            self._current_cpu_bytes = 0

    def __del__(self):
        try:
            self.shutdown()
        except Exception:
            pass
