from __future__ import annotations

from types import SimpleNamespace

import pytest
import torch.nn as nn

from obliteratus.models import loader
from obliteratus.models.loader import ModelHandle
from obliteratus.strategies.utils import (
    get_attention_module,
    get_embedding_module,
    get_ffn_module,
    get_layer_modules,
)


class _Gemma4LikeLayer(nn.Module):
    def __init__(self):
        super().__init__()
        self.self_attn = nn.Module()
        self.mlp = nn.Module()


class _Gemma4LikeModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.model = nn.Module()
        self.model.language_model = nn.Module()
        self.model.language_model.layers = nn.ModuleList(
            [_Gemma4LikeLayer(), _Gemma4LikeLayer()]
        )
        self.model.language_model.embed_tokens = nn.Embedding(16, 8)
        self.lm_head = nn.Linear(8, 16, bias=False)


def _gemma4_config(model_type: str = "gemma4_unified"):
    return SimpleNamespace(
        model_type=model_type,
        architectures=["Gemma4UnifiedForConditionalGeneration"],
        text_config=SimpleNamespace(
            num_hidden_layers=2,
            num_attention_heads=4,
            hidden_size=8,
            intermediate_size=32,
        ),
    )


def test_gemma4_unified_uses_image_text_auto_model_when_available():
    if loader.AutoModelForImageTextToText is None:
        with pytest.raises(RuntimeError, match="AutoModelForImageTextToText"):
            loader._select_model_class("causal_lm", _gemma4_config())
        return

    assert (
        loader._select_model_class("causal_lm", _gemma4_config())
        is loader.AutoModelForImageTextToText
    )


def test_gemma4_nested_language_stack_paths():
    handle = ModelHandle(
        model=_Gemma4LikeModel(),
        tokenizer=SimpleNamespace(pad_token="<pad>", eos_token="<eos>"),
        config=_gemma4_config(),
        model_name="google/gemma-4-12B-it",
        task="causal_lm",
    )

    layers = get_layer_modules(handle)

    assert len(layers) == 2
    assert get_attention_module(layers[0], handle.architecture) is layers[0].self_attn
    assert get_ffn_module(layers[0], handle.architecture) is layers[0].mlp
    assert get_embedding_module(handle).num_embeddings == 16
    assert handle.num_layers == 2
    assert handle.hidden_size == 8