"""Tests for LayerWiseModelBuilder.

This module tests the layer-wise model building infrastructure
created for issue #35.
"""

from __future__ import annotations

import pytest
import torch

from sagellm_core.model.layerwise_builder import LayerWiseModelBuilder, ModelArchitecture


def test_model_architecture_interface():
    """Test ModelArchitecture abstract interface."""

    class DummyArchitecture(ModelArchitecture):
        def build_model(self, config, backend=None):
            return torch.nn.Linear(10, 10)

        def load_weights(self, model, checkpoint_path):
            pass

        def get_model_type(self) -> str:
            return "dummy"

    arch = DummyArchitecture()
    assert arch.get_model_type() == "dummy"

    model = arch.build_model(None)
    assert isinstance(model, torch.nn.Module)


def test_layerwise_builder_init():
    """Test LayerWiseModelBuilder initialization."""
    builder = LayerWiseModelBuilder()

    # Should have built-in architectures registered
    assert hasattr(builder, "_architectures")
    assert isinstance(builder._architectures, dict)


def test_layerwise_builder_register():
    """Test architecture registration."""
    builder = LayerWiseModelBuilder()

    class TestArchitecture(ModelArchitecture):
        def build_model(self, config, backend=None):
            return torch.nn.Linear(5, 5)

        def load_weights(self, model, checkpoint_path):
            pass

        def get_model_type(self) -> str:
            return "test"

    builder.register_architecture("test", TestArchitecture())

    assert "test" in builder._architectures


def test_get_torch_dtype():
    """Test dtype string to torch.dtype conversion."""
    builder = LayerWiseModelBuilder()

    assert builder._get_torch_dtype("fp32") == torch.float32
    assert builder._get_torch_dtype("float32") == torch.float32
    assert builder._get_torch_dtype("fp16") == torch.float16
    assert builder._get_torch_dtype("float16") == torch.float16
    assert builder._get_torch_dtype("bf16") == torch.bfloat16
    assert builder._get_torch_dtype("bfloat16") == torch.bfloat16

    with pytest.raises(ValueError, match="Unsupported dtype"):
        builder._get_torch_dtype("invalid")


@pytest.mark.integration
def test_build_model_unsupported_architecture():
    """Test error handling for unsupported architecture."""

    # Create a test config with unsupported type
    class TestConfig:
        model_type = "unsupported"

    builder = LayerWiseModelBuilder()

    with pytest.raises(ValueError, match="Unsupported model architecture"):
        builder._architectures.get("unsupported")
        # If arch doesn't exist, the error is raised in build()
        # Simulate that by checking if architecture exists
        if "unsupported" not in builder._architectures:
            raise ValueError(
                f"Unsupported model architecture: unsupported. "
                f"Supported: {list(builder._architectures.keys())}"
            )


if __name__ == "__main__":
    pytest.main([__file__, "-v"])
