"""
Tests for the Model Graph Compiler.

These tests verify the model-agnostic structure discovery using
synthetic toy models rather than real multi-GB models.
"""

import json
import os
import tempfile

import pytest
import torch
import torch.nn as nn

from dokodemo_ai.graph.compiler import (
    ModelGraphCompiler,
    ModelGraph,
    LayerInfo,
    EmbeddingInfo,
    HeadInfo,
    TensorRef,
)


# ---------------------------------------------------------------------------
# Toy model fixtures
# ---------------------------------------------------------------------------


class ToyMLP(nn.Module):
    """Simple two-layer MLP (one 'expert' in isolation)."""

    def __init__(self, d_model: int, d_ffn: int):
        super().__init__()
        self.gate_proj = nn.Linear(d_model, d_ffn, bias=False)
        self.up_proj = nn.Linear(d_model, d_ffn, bias=False)
        self.down_proj = nn.Linear(d_ffn, d_model, bias=False)

    def forward(self, x):
        return self.down_proj(
            torch.nn.functional.silu(self.gate_proj(x)) * self.up_proj(x)
        )


class ToyAttention(nn.Module):
    def __init__(self, d_model: int, num_heads: int):
        super().__init__()
        self.q_proj = nn.Linear(d_model, d_model, bias=False)
        self.k_proj = nn.Linear(d_model, d_model, bias=False)
        self.v_proj = nn.Linear(d_model, d_model, bias=False)
        self.o_proj = nn.Linear(d_model, d_model, bias=False)

    def forward(self, x):
        return self.o_proj(self.q_proj(x))


class ToyLayer(nn.Module):
    def __init__(self, d_model: int, d_ffn: int, num_heads: int):
        super().__init__()
        self.input_layernorm = nn.RMSNorm(d_model, eps=1e-6)
        self.self_attn = ToyAttention(d_model, num_heads)
        self.post_attention_layernorm = nn.RMSNorm(d_model, eps=1e-6)
        self.mlp = ToyMLP(d_model, d_ffn)

    def forward(self, x):
        x = x + self.self_attn(self.input_layernorm(x))
        x = x + self.mlp(self.post_attention_layernorm(x))
        return x


class ToyTransformer(nn.Module):
    def __init__(
        self,
        vocab_size: int = 1000,
        d_model: int = 64,
        d_ffn: int = 128,
        num_layers: int = 4,
        num_heads: int = 4,
    ):
        super().__init__()
        self.embed_tokens = nn.Embedding(vocab_size, d_model)
        self.layers = nn.ModuleList([
            ToyLayer(d_model, d_ffn, num_heads)
            for _ in range(num_layers)
        ])
        self.norm = nn.RMSNorm(d_model, eps=1e-6)
        self.lm_head = nn.Linear(d_model, vocab_size, bias=False)

    def forward(self, x):
        h = self.embed_tokens(x)
        for layer in self.layers:
            h = layer(h)
        return self.lm_head(self.norm(h))


class ToyMoERouter(nn.Module):
    def __init__(self, d_model: int, num_experts: int):
        super().__init__()
        self.gate = nn.Linear(d_model, num_experts, bias=False)

    def forward(self, x):
        return self.gate(x)


class ToyMoEBlock(nn.Module):
    def __init__(self, d_model: int, d_ffn: int, num_experts: int):
        super().__init__()
        self.gate = ToyMoERouter(d_model, num_experts)
        self.experts = nn.ModuleList([
            ToyMLP(d_model, d_ffn) for _ in range(num_experts)
        ])

    def forward(self, x):
        return x


class ToyMoELayer(nn.Module):
    def __init__(self, d_model: int, d_ffn: int, num_experts: int, num_heads: int):
        super().__init__()
        self.input_layernorm = nn.RMSNorm(d_model, eps=1e-6)
        self.self_attn = ToyAttention(d_model, num_heads)
        self.post_attention_layernorm = nn.RMSNorm(d_model, eps=1e-6)
        self.block_sparse_moe = ToyMoEBlock(d_model, d_ffn, num_experts)


class ToyMoETransformer(nn.Module):
    def __init__(
        self,
        vocab_size: int = 1000,
        d_model: int = 64,
        d_ffn: int = 128,
        num_layers: int = 4,
        num_heads: int = 4,
        num_experts: int = 8,
    ):
        super().__init__()
        self.embed_tokens = nn.Embedding(vocab_size, d_model)
        self.layers = nn.ModuleList([
            ToyMoELayer(d_model, d_ffn, num_experts, num_heads)
            for _ in range(num_layers)
        ])
        self.norm = nn.RMSNorm(d_model, eps=1e-6)
        self.lm_head = nn.Linear(d_model, vocab_size, bias=False)


# ---------------------------------------------------------------------------
# Helpers to write a toy model to a temp directory as SafeTensors
# ---------------------------------------------------------------------------


def _save_toy_model(model: nn.Module, config: dict, tmp_dir: str) -> None:
    """Save a toy model to SafeTensors format in tmp_dir."""
    from safetensors.torch import save_file

    params = {name: param.data for name, param in model.named_parameters()}

    # Write as a single file
    save_file(params, os.path.join(tmp_dir, "model.safetensors"))

    # Write config
    with open(os.path.join(tmp_dir, "config.json"), "w") as f:
        json.dump(config, f)


def _toy_dense_config(
    vocab_size=1000, d_model=64, d_ffn=128, num_layers=4, num_heads=4
):
    return {
        "model_type": "toy_dense",
        "vocab_size": vocab_size,
        "hidden_size": d_model,
        "intermediate_size": d_ffn,
        "num_hidden_layers": num_layers,
        "num_attention_heads": num_heads,
        "num_key_value_heads": num_heads,
        "rms_norm_eps": 1e-6,
    }


def _toy_moe_config(
    vocab_size=1000, d_model=64, d_ffn=128, num_layers=4,
    num_heads=4, num_experts=8
):
    cfg = _toy_dense_config(vocab_size, d_model, d_ffn, num_layers, num_heads)
    cfg.update({
        "model_type": "toy_moe",
        "num_local_experts": num_experts,
        "num_experts_per_tok": 2,
    })
    return cfg


# ---------------------------------------------------------------------------
# Tests for TensorRef and data structures
# ---------------------------------------------------------------------------


class TestTensorRef:
    def test_size_bytes_fp16(self):
        ref = TensorRef(name="foo", file="f.st", dtype="F16", shape=[1024, 1024])
        assert ref.numel == 1024 * 1024
        assert ref.size_bytes == 1024 * 1024 * 2  # 2 bytes per fp16 element

    def test_size_bytes_f32(self):
        ref = TensorRef(name="foo", file="f.st", dtype="F32", shape=[512, 512])
        assert ref.size_bytes == 512 * 512 * 4

    def test_size_bytes_int8(self):
        ref = TensorRef(name="foo", file="f.st", dtype="I8", shape=[256])
        assert ref.size_bytes == 256


# ---------------------------------------------------------------------------
# Tests for Layer discovery
# ---------------------------------------------------------------------------


class TestLayerDiscovery:
    """Test that the compiler correctly finds layers in toy models."""

    def test_find_layer_list_dense(self):
        """Should find the transformer layer list."""
        model = ToyTransformer(num_layers=4)
        compiler = ModelGraphCompiler()
        layers = compiler._find_layer_list(model)
        assert len(layers) == 4

    def test_layer_names_are_indexed(self):
        """Layer module names should be indexed correctly."""
        model = ToyTransformer(num_layers=3)
        compiler = ModelGraphCompiler()
        layers = compiler._find_layer_list(model)
        assert "layers.0" in layers[0][0]
        assert "layers.1" in layers[1][0]
        assert "layers.2" in layers[2][0]

    def test_find_embedding(self):
        """Should locate the embedding layer."""
        model = ToyTransformer()
        compiler = ModelGraphCompiler()
        # Fake param_meta mapping
        param_meta = {
            name: TensorRef(name=name, file="f.st", dtype="F16", shape=[1])
            for name, _ in model.named_parameters()
        }
        emb = compiler._find_embedding(model, param_meta)
        assert emb is not None
        assert "embed_tokens" in emb.module_name

    def test_no_layer_list_raises(self):
        """A model without nn.ModuleList should raise RuntimeError."""
        class NoList(nn.Module):
            def __init__(self):
                super().__init__()
                self.linear = nn.Linear(64, 64)

        compiler = ModelGraphCompiler()
        with pytest.raises(RuntimeError, match="layer list"):
            compiler._find_layer_list(NoList())


# ---------------------------------------------------------------------------
# Tests for MoE detection
# ---------------------------------------------------------------------------


class TestMoEDetection:
    def test_dense_layer_not_moe(self):
        """ToyLayer should be classified as dense (not MoE)."""
        layer = ToyLayer(d_model=64, d_ffn=128, num_heads=4)
        compiler = ModelGraphCompiler()
        assert not compiler._has_expert_children(layer)

    def test_moe_block_detected(self):
        """ToyMoEBlock should be detected as containing experts."""
        moe_block = ToyMoEBlock(d_model=64, d_ffn=128, num_experts=8)
        compiler = ModelGraphCompiler()
        assert compiler._has_expert_children(moe_block)

    def test_moe_layer_analysis(self):
        """MoE layer should produce LayerInfo with is_moe=True."""
        model = ToyMoETransformer(num_experts=8)
        compiler = ModelGraphCompiler()
        param_meta = {
            name: TensorRef(name=name, file="f.st", dtype="F16", shape=[1])
            for name, _ in model.named_parameters()
        }
        layer_modules = compiler._find_layer_list(model)
        config = _toy_moe_config(num_experts=8)
        layers = compiler._analyze_layers(layer_modules, param_meta, config)

        assert all(layer.is_moe for layer in layers)
        assert all(layer.num_experts == 8 for layer in layers)
        assert all(layer.num_active_experts == 2 for layer in layers)

    def test_experts_discoverable(self):
        """Expert components should be individually addressable."""
        model = ToyMoETransformer(num_experts=4)
        compiler = ModelGraphCompiler()
        param_meta = {
            name: TensorRef(name=name, file="f.st", dtype="F16", shape=[1])
            for name, _ in model.named_parameters()
        }
        layer_modules = compiler._find_layer_list(model)
        config = _toy_moe_config(num_experts=4)
        layers = compiler._analyze_layers(layer_modules, param_meta, config)

        # Each layer should have individual expert components
        for layer in layers:
            experts = layer.experts
            assert len(experts) == 4, (
                f"Expected 4 experts, got {len(experts)}"
            )


# ---------------------------------------------------------------------------
# Tests for full compile pipeline (with temp SafeTensors files)
# ---------------------------------------------------------------------------


class TestFullCompile:
    """Integration tests using real SafeTensors files."""

    def test_dense_model_compile(self, tmp_path):
        """Full compile pipeline on a toy dense model."""
        pytest.importorskip("safetensors")
        pytest.importorskip("transformers")
        pytest.importorskip("accelerate")

        model = ToyTransformer(num_layers=2, d_model=32, d_ffn=64)
        config = _toy_dense_config(d_model=32, d_ffn=64, num_layers=2)

        _save_toy_model(model, config, str(tmp_path))

        # The compiler uses AutoModelForCausalLM which needs a real
        # registered model_type; we test the SafeTensors index loading instead.
        from dokodemo_ai.graph.compiler import ModelGraphCompiler
        compiler = ModelGraphCompiler()

        # Test index loading
        index = compiler._load_safetensors_index(tmp_path)
        assert len(index) > 0

        # Verify all params are referenced
        all_params = {n for n, _ in model.named_parameters()}
        for param_name in all_params:
            assert param_name in index, (
                f"Missing param in index: {param_name}"
            )

    def test_safetensors_index_shape_extraction(self, tmp_path):
        """Tensor shapes should be extractable from SafeTensors without loading."""
        pytest.importorskip("safetensors")

        model = ToyTransformer(num_layers=1, d_model=16, d_ffn=32)
        config = _toy_dense_config(d_model=16, d_ffn=32, num_layers=1)
        _save_toy_model(model, config, str(tmp_path))

        compiler = ModelGraphCompiler()
        index = compiler._load_safetensors_index(tmp_path)
        meta = compiler._build_param_metadata(tmp_path, index)

        # Check a known parameter
        assert "embed_tokens.weight" in meta
        ref = meta["embed_tokens.weight"]
        assert ref.shape == [1000, 16]  # [vocab_size, d_model]
        assert ref.dtype in ("F32", "F16", "BF16")
