"""Tests for Model Loader (HF-first architecture).

This test module validates:
1. ModelLoadConfig creation and conversion
2. Loader factory always returns HFModelLoader
3. HFModelLoader backward compatibility
4. LLMEngine integration

Architecture decision (2026-02):
    sageLLM uses HuggingFace models exclusively. Custom model reimplementations
    have been removed. All models are loaded via HFModelLoader.
"""

import pytest
import torch

from sagellm_core.llm_engine import LLMEngine, LLMEngineConfig
from sagellm_core.model import (
    HFModelLoader,
    ModelLoadConfig,
    get_model_loader,
)


class TestModelLoadConfig:
    """Test ModelLoadConfig creation and validation."""

    def test_create_basic_config(self):
        """Test creating basic ModelLoadConfig."""
        config = ModelLoadConfig(
            model_path="gpt2",
            dtype="float16",
        )

        assert config.model_path == "gpt2"
        assert config.dtype == "float16"
        assert config.tensor_parallel_size == 1

    def test_from_engine_config(self):
        """Test creating ModelLoadConfig from LLMEngineConfig."""
        engine_config = LLMEngineConfig(
            model_path="gpt2",
            dtype="float32",
            max_model_len=2048,
            tensor_parallel_size=1,
        )

        load_config = ModelLoadConfig.from_engine_config(engine_config, backend=None)

        assert load_config.model_path == "gpt2"
        assert load_config.dtype == "float32"
        assert load_config.max_model_len == 2048
        assert load_config.tensor_parallel_size == 1

    def test_get_torch_dtype(self):
        """Test torch dtype resolution."""
        config = ModelLoadConfig(model_path="gpt2", dtype="float16")
        assert config.get_torch_dtype() == torch.float16

        config = ModelLoadConfig(model_path="gpt2", dtype="bfloat16")
        assert config.get_torch_dtype() == torch.bfloat16

        config = ModelLoadConfig(model_path="gpt2", dtype="float32")
        assert config.get_torch_dtype() == torch.float32

    def test_get_torch_dtype_auto(self):
        """Test auto dtype selection."""
        config = ModelLoadConfig(model_path="gpt2", dtype="auto", device="cpu")
        assert config.get_torch_dtype() == torch.float32

        config = ModelLoadConfig(model_path="gpt2", dtype="auto", device="cuda")
        assert config.get_torch_dtype() == torch.float16

    def test_get_device_map(self):
        """Test device map generation for HF transformers."""
        config = ModelLoadConfig(model_path="gpt2", device="cpu")
        assert config.get_device_map() == "cpu"

        config = ModelLoadConfig(model_path="gpt2", device="cuda")
        assert config.get_device_map() == "auto"

        config = ModelLoadConfig(model_path="gpt2", device="auto")
        assert config.get_device_map() == "auto"


class TestLoaderFactory:
    """Test loader selection factory (get_model_loader)."""

    def test_always_returns_hf_loader(self):
        """Test that get_model_loader always returns HFModelLoader."""
        config = ModelLoadConfig(model_path="gpt2")
        loader = get_model_loader(config)
        assert isinstance(loader, HFModelLoader)

    def test_hf_loader_for_any_model(self):
        """Test HFModelLoader for any model path."""
        config = ModelLoadConfig(model_path="meta-llama/Llama-2-7b-hf")
        loader = get_model_loader(config)
        assert isinstance(loader, HFModelLoader)

    def test_hf_loader_for_local_path(self):
        """Test HFModelLoader for local model path."""
        config = ModelLoadConfig(model_path="/tmp/some-model")
        loader = get_model_loader(config)
        assert isinstance(loader, HFModelLoader)


class TestHFModelLoader:
    """Test HFModelLoader functionality."""

    @pytest.mark.slow
    def test_load_tiny_model(self):
        """Test loading tiny-gpt2 with HFModelLoader."""
        config = ModelLoadConfig(
            model_path="sshleifer/tiny-gpt2",
            dtype="float32",
            device="cpu",
        )

        loader = HFModelLoader()
        model, tokenizer = loader.load(config)

        assert model is not None
        assert tokenizer is not None
        assert hasattr(model, "forward")
        assert hasattr(tokenizer, "encode")

    def test_hf_loader_supports_all(self):
        """Test HFModelLoader.supports() returns True for all configs."""
        config = ModelLoadConfig(model_path="any-model")
        loader = HFModelLoader()

        assert loader.supports(config) is True


class TestLLMEngineIntegration:
    """Test LLMEngine integration with HFModelLoader."""

    @pytest.mark.slow
    @pytest.mark.asyncio
    async def test_llm_engine_with_hf_loader(self):
        """Test LLMEngine uses HFModelLoader."""
        config = LLMEngineConfig(
            model_path="sshleifer/tiny-gpt2",
            backend_type="cpu",
            max_batch_size=1,
        )

        engine = LLMEngine(config)
        await engine.start()

        assert engine.is_running
        assert engine._model is not None
        assert engine._tokenizer is not None

        # Test generation
        response = await engine.generate("Hello", max_tokens=5)
        assert response.output_text

        await engine.stop()

    @pytest.mark.asyncio
    async def test_llm_engine_config_conversion(self):
        """Test LLMEngine creates correct ModelLoadConfig."""
        config = LLMEngineConfig(
            model_path="sshleifer/tiny-gpt2",
            backend_type="cpu",
            dtype="float32",
            max_model_len=512,
        )

        engine = LLMEngine(config)
        await engine.start()

        # Verify model was loaded
        assert engine._model is not None

        await engine.stop()


class TestBackwardCompatibility:
    """Test backward compatibility with existing LLMEngine behavior."""

    @pytest.mark.slow
    @pytest.mark.asyncio
    async def test_existing_code_still_works(self):
        """Test existing LLMEngine usage patterns still work."""
        config = LLMEngineConfig(
            model_path="sshleifer/tiny-gpt2",
            backend_type="cpu",
        )

        engine = LLMEngine(config)
        await engine.start()

        response = await engine.generate(
            prompt="Hello world",
            max_tokens=10,
            temperature=0.0,
        )

        assert response.output_text
        assert response.metrics.ttft_ms >= 0
        assert response.metrics.throughput_tps >= 0

        await engine.stop()

    @pytest.mark.slow
    @pytest.mark.asyncio
    async def test_sampling_params_still_work(self):
        """Test sampling parameter handling still works."""
        from sagellm_protocol.sampling import DecodingStrategy, SamplingParams

        config = LLMEngineConfig(
            model_path="sshleifer/tiny-gpt2",
            backend_type="cpu",
        )

        engine = LLMEngine(config)
        await engine.start()

        # Test with SamplingParams object
        sampling_params = SamplingParams(
            strategy=DecodingStrategy.GREEDY,
            max_tokens=5,
        )

        response = await engine.generate(
            prompt="Hello",
            sampling_params=sampling_params,
        )

        assert response.output_text

        await engine.stop()


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