"""
Tests for Adaptive Layer Cache.

Verifies:
1. Basic get/put operations
2. LRU eviction with importance weighting
3. Memory budget enforcement
4. Expert-specific eviction for MoE layers
"""

import pytest
import torch

from dokodemo_ai.engine.cache import AdaptiveLayerCache, CacheEntry
from dokodemo_ai.graph.compiler import ComponentInfo, LayerInfo, ModelGraph, TensorRef


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------


def _make_graph(num_layers: int = 8) -> ModelGraph:
    """Create a minimal ModelGraph for testing."""
    from dokodemo_ai.graph.compiler import EmbeddingInfo, HeadInfo

    layers = [
        LayerInfo(index=i, module_name=f"model.layers.{i}")
        for i in range(num_layers)
    ]
    return ModelGraph(
        model_path="/fake",
        config={"hidden_size": 64, "vocab_size": 1000},
        embedding=EmbeddingInfo(module_name="embed_tokens"),
        layers=layers,
        head=HeadInfo(module_name="lm_head"),
        safetensors_index={},
        num_layers=num_layers,
        hidden_size=64,
        vocab_size=1000,
    )


def _make_component(name: str, kind: str = "attention", num_floats: int = 1024) -> ComponentInfo:
    ref = TensorRef(name=f"{name}.weight", file="model.safetensors", dtype="F16", shape=[num_floats])
    return ComponentInfo(name=name, kind=kind, params=[ref])


def _make_tensors(size_floats: int = 1024) -> dict[str, torch.Tensor]:
    t = torch.randn(size_floats, dtype=torch.float16)
    return {"weight": t}


# ---------------------------------------------------------------------------
# Tests
# ---------------------------------------------------------------------------


class TestCacheBasicOps:
    def test_put_and_get(self):
        device = torch.device("cpu")
        graph = _make_graph(4)
        cache = AdaptiveLayerCache(
            memory_budget_bytes=10 * 1024**2,  # 10 MB
            model_graph=graph,
            device=device,
        )
        comp = _make_component("model.layers.0.self_attn")
        tensors = _make_tensors(512)

        cache.put(comp, tensors, layer_index=0)
        result = cache.get(comp)

        assert result is not None
        assert "weight" in result

    def test_miss_returns_none(self):
        device = torch.device("cpu")
        graph = _make_graph(4)
        cache = AdaptiveLayerCache(
            memory_budget_bytes=10 * 1024**2,
            model_graph=graph,
            device=device,
        )
        comp = _make_component("nonexistent.layer")
        assert cache.get(comp) is None

    def test_put_same_key_twice(self):
        device = torch.device("cpu")
        graph = _make_graph(4)
        cache = AdaptiveLayerCache(
            memory_budget_bytes=10 * 1024**2,
            model_graph=graph,
            device=device,
        )
        comp = _make_component("model.layers.0.self_attn")
        tensors = _make_tensors()

        cache.put(comp, tensors, layer_index=0)
        # Second put should succeed (already cached)
        result = cache.put(comp, tensors, layer_index=0)
        assert result is True

    def test_explicit_eviction(self):
        device = torch.device("cpu")
        graph = _make_graph(4)
        cache = AdaptiveLayerCache(
            memory_budget_bytes=10 * 1024**2,
            model_graph=graph,
            device=device,
        )
        comp = _make_component("model.layers.1.mlp")
        cache.put(comp, _make_tensors(), layer_index=1)
        cache.evict_component(comp.name)
        assert cache.get(comp) is None

    def test_clear(self):
        device = torch.device("cpu")
        graph = _make_graph(4)
        cache = AdaptiveLayerCache(
            memory_budget_bytes=10 * 1024**2,
            model_graph=graph,
            device=device,
        )
        for i in range(4):
            comp = _make_component(f"model.layers.{i}.self_attn")
            cache.put(comp, _make_tensors(), layer_index=i)

        cache.clear()
        assert cache.current_bytes == 0
        assert len(cache._cache) == 0


class TestCacheMemoryBudget:
    def test_budget_enforced(self):
        """Cache should not exceed memory budget."""
        device = torch.device("cpu")
        graph = _make_graph(8)
        budget = 8 * 1024  # 8 KB — very tight

        cache = AdaptiveLayerCache(
            memory_budget_bytes=budget,
            model_graph=graph,
            device=device,
        )

        # Each tensor is 2KB (1024 float16 elements)
        for i in range(8):
            comp = _make_component(f"model.layers.{i}.mlp", num_floats=1024)
            tensors = _make_tensors(1024)  # 2 KB per tensor
            cache.put(comp, tensors, layer_index=i)

        assert cache.current_bytes <= budget, (
            f"Cache ({cache.current_bytes} bytes) exceeds budget ({budget} bytes)"
        )

    def test_oversized_component_not_cached(self):
        """A component larger than the entire budget should be rejected."""
        device = torch.device("cpu")
        graph = _make_graph(4)
        budget = 100  # 100 bytes — tiny

        cache = AdaptiveLayerCache(
            memory_budget_bytes=budget,
            model_graph=graph,
            device=device,
        )

        comp = _make_component("large.layer", num_floats=10000)
        tensors = _make_tensors(10000)  # ~20 KB
        result = cache.put(comp, tensors, layer_index=0)
        assert result is False


class TestCacheImportanceScoring:
    def test_importance_scores_computed(self):
        """All layers should have importance scores."""
        device = torch.device("cpu")
        n = 10
        graph = _make_graph(n)
        cache = AdaptiveLayerCache(
            memory_budget_bytes=100 * 1024**2,
            model_graph=graph,
            device=device,
        )
        scores = cache._importance_scores
        assert len(scores) == n
        for i in range(n):
            assert i in scores
            assert 0.0 < scores[i] <= 1.0

    def test_edge_layers_more_important(self):
        """First and last layers should have higher importance than middle."""
        device = torch.device("cpu")
        n = 10
        graph = _make_graph(n)
        cache = AdaptiveLayerCache(
            memory_budget_bytes=100 * 1024**2,
            model_graph=graph,
            device=device,
        )
        scores = cache._importance_scores
        assert scores[0] > scores[n // 2], "First layer should be most important"
        assert scores[n - 1] > scores[n // 2], "Last layer should be more important than middle"

    def test_router_gets_importance_boost(self):
        """Router components should get higher cached importance."""
        device = torch.device("cpu")
        graph = _make_graph(4)
        cache = AdaptiveLayerCache(
            memory_budget_bytes=10 * 1024**2,
            model_graph=graph,
            device=device,
        )

        router_comp = _make_component("layer.0.router", kind="router")
        mlp_comp = _make_component("layer.0.mlp", kind="mlp")

        cache.put(router_comp, _make_tensors(256), layer_index=0)
        cache.put(mlp_comp, _make_tensors(256), layer_index=0)

        router_entry = cache._cache.get(router_comp.name)
        mlp_entry = cache._cache.get(mlp_comp.name)

        assert router_entry is not None
        assert mlp_entry is not None
        assert router_entry.importance > mlp_entry.importance


class TestCacheStats:
    def test_stats_returns_dict(self):
        device = torch.device("cpu")
        graph = _make_graph(4)
        cache = AdaptiveLayerCache(
            memory_budget_bytes=10 * 1024**2,
            model_graph=graph,
            device=device,
        )
        stats = cache.stats
        assert "entries" in stats
        assert "used_bytes" in stats
        assert "budget_bytes" in stats
        assert "utilization" in stats
        assert stats["utilization"] == 0.0  # Empty cache

    def test_stats_after_put(self):
        device = torch.device("cpu")
        graph = _make_graph(4)
        cache = AdaptiveLayerCache(
            memory_budget_bytes=10 * 1024**2,
            model_graph=graph,
            device=device,
        )
        comp = _make_component("layer.0.attn")
        cache.put(comp, _make_tensors(512), layer_index=0)

        stats = cache.stats
        assert stats["entries"] == 1
        assert stats["used_bytes"] > 0
        assert stats["utilization"] > 0.0
