"""
Adaptive Layer Cache - Importance-weighted LRU caching for GPU memory.

Unlike AirLLM (which evicts every layer after use), Dokodemo AI maintains
an intelligent cache of layer weights in GPU memory. Layers are scored by
importance and the cache retains the most valuable layers across tokens.

Importance scoring factors:
1. Position: First and last layers are most important for output quality
2. Frequency: Layers accessed more recently are prioritized
3. Size efficiency: Smaller layers are cheaper to keep resident
4. MoE structure: Non-expert components are always cached if possible
   (since they're needed for every token regardless of routing)
"""

from __future__ import annotations

import logging
import threading
from collections import OrderedDict
from dataclasses import dataclass
from typing import Optional

import torch

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

logger = logging.getLogger(__name__)


@dataclass
class CacheEntry:
    """A cached component in GPU memory."""

    component_name: str
    kind: str
    tensors: dict[str, torch.Tensor]
    size_bytes: int
    importance: float
    access_count: int = 0
    layer_index: int = 0


class AdaptiveLayerCache:
    """
    GPU-resident cache for model layer weights.

    Manages a fixed GPU memory budget, keeping the most important
    layers/components resident to avoid redundant disk I/O.
    """

    def __init__(
        self,
        memory_budget_bytes: int,
        model_graph: ModelGraph,
        device: torch.device,
    ):
        """
        Args:
            memory_budget_bytes: GPU memory budget for caching weights.
            model_graph: Compiled model graph for importance scoring.
            device: GPU device.
        """
        self.budget = memory_budget_bytes
        self.graph = model_graph
        self.device = device
        self.current_bytes = 0

        # Cache storage: component_name -> CacheEntry
        self._cache: OrderedDict[str, CacheEntry] = OrderedDict()
        self._lock = threading.Lock()

        # Pre-compute importance scores for all layers
        self._importance_scores = self._compute_importance_scores()

        logger.info(
            "AdaptiveLayerCache initialized: %.2f GB budget on %s",
            memory_budget_bytes / 1e9,
            device,
        )

    def _compute_importance_scores(self) -> dict[int, float]:
        """
        Pre-compute importance scores for each layer.

        Uses a U-shaped curve: first and last layers are most important.
        This is grounded in research showing that early layers handle
        tokenization/embedding and late layers handle output projection,
        both of which are critical for quality.
        """
        n = self.graph.num_layers
        scores: dict[int, float] = {}

        for i in range(n):
            # U-shaped importance: 1.0 at edges, 0.2 at center
            distance_from_edge = min(i, n - 1 - i)
            normalized = distance_from_edge / max(n // 2, 1)
            score = 1.0 - 0.8 * normalized
            scores[i] = score

        return scores

    def get(self, component: ComponentInfo) -> Optional[dict[str, torch.Tensor]]:
        """
        Get a cached component's tensors, or None if not cached.

        Updates access statistics for the LRU policy.
        """
        with self._lock:
            entry = self._cache.get(component.name)
            if entry is None:
                return None

            # Update access stats
            entry.access_count += 1
            self._cache.move_to_end(component.name)
            return entry.tensors

    def put(
        self,
        component: ComponentInfo,
        tensors: dict[str, torch.Tensor],
        layer_index: int,
    ) -> bool:
        """
        Cache a component's tensors in GPU memory.

        Returns True if cached successfully, False if eviction couldn't
        free enough space (component is too large for budget).
        """
        size = sum(t.nbytes for t in tensors.values())

        if size > self.budget:
            return False  # Single component exceeds entire budget

        with self._lock:
            # Already cached?
            if component.name in self._cache:
                self._cache.move_to_end(component.name)
                return True

            # Make room
            self._evict_until(self.budget - size)

            importance = self._importance_scores.get(layer_index, 0.5)

            # Boost importance for non-expert MoE components
            # (routers and attention are needed for every token)
            if component.kind in ("router", "attention", "norm"):
                importance *= 1.5

            entry = CacheEntry(
                component_name=component.name,
                kind=component.kind,
                tensors=tensors,
                size_bytes=size,
                importance=importance,
                layer_index=layer_index,
            )

            self._cache[component.name] = entry
            self.current_bytes += size
            return True

    def get_or_load(
        self,
        component: ComponentInfo,
        layer_index: int,
        loader,
    ) -> dict[str, torch.Tensor]:
        """
        Get from cache or load using the provided loader function.

        Args:
            component: Component to get.
            layer_index: Layer index for importance scoring.
            loader: Callable that returns dict[str, torch.Tensor].

        Returns:
            Component tensors on GPU.
        """
        cached = self.get(component)
        if cached is not None:
            return cached

        tensors = loader()
        self.put(component, tensors, layer_index)
        return tensors

    def evict_component(self, component_name: str) -> None:
        """Explicitly evict a component."""
        with self._lock:
            entry = self._cache.pop(component_name, None)
            if entry is not None:
                self.current_bytes -= entry.size_bytes
                # Free GPU memory
                for t in entry.tensors.values():
                    del t
                entry.tensors.clear()

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

    def evict_experts(self, layer: LayerInfo) -> None:
        """Evict only expert components of a MoE layer."""
        for comp in layer.components:
            if comp.kind == "expert":
                self.evict_component(comp.name)

    def _evict_until(self, target_bytes: int) -> None:
        """Evict lowest-importance entries until current_bytes <= target."""
        while self.current_bytes > target_bytes and self._cache:
            # Find the entry with lowest effective importance
            # Effective importance = base importance * recency boost
            worst_key = None
            worst_score = float("inf")

            for key, entry in self._cache.items():
                # Recency boost: recently accessed items get a bonus
                recency = 1.0 + 0.1 * entry.access_count
                effective = entry.importance * recency

                # Size penalty: prefer evicting large items to free more space
                size_factor = entry.size_bytes / (1024**2)  # MB
                score = effective / max(size_factor, 0.1)

                if score < worst_score:
                    worst_score = score
                    worst_key = key

            if worst_key is None:
                break

            entry = self._cache.pop(worst_key)
            self.current_bytes -= entry.size_bytes
            for t in entry.tensors.values():
                del t
            entry.tensors.clear()

    def clear(self) -> None:
        """Clear all cached entries."""
        with self._lock:
            for entry in self._cache.values():
                for t in entry.tensors.values():
                    del t
                entry.tensors.clear()
            self._cache.clear()
            self.current_bytes = 0

    @property
    def stats(self) -> dict:
        """Return cache statistics."""
        with self._lock:
            return {
                "entries": len(self._cache),
                "used_bytes": self.current_bytes,
                "budget_bytes": self.budget,
                "utilization": self.current_bytes / max(self.budget, 1),
                "cached_components": list(self._cache.keys()),
            }
