"""Integration tests for LLM2SLM.
This module contains integration tests that test the full workflow
from model conversion to export, loading, and inference.
"""

import tempfile
from pathlib import Path

import pytest

from llm2slm.core.config import Config
from llm2slm.slm.export import SLMExporter
from llm2slm.slm.runtime import SLMRuntime


class TestIntegration:
    """Integration tests for the full LLM2SLM workflow."""

    @pytest.mark.asyncio
    async def test_full_conversion_workflow(self):
        """Test the complete workflow: convert -> export -> load -> infer."""
        config = Config(openai_api_key="test_key")
        # Step 1: Simulate model conversion (mock the conversion process)
        model_data = {
            "parameters": 500000,
            "size": 250000,
            "compression_ratio": 0.5,
            "quantization": "8bit",
            "pruned": True,
            "distilled": False,
            "architecture": "transformer",
            "layers": ["layer1", "layer2"],
            "hidden_size": 768,
            "attention_heads": 12,
            "tokenizer": {"vocab_size": 30000},
        }
        model_info = {
            "model_id": "gpt-3.5-turbo",
            "provider": "openai",
            "parameters": 175000000,
            "size": 1000000,
        }
        # Step 2: Export the model
        exporter = SLMExporter(config)
        with tempfile.TemporaryDirectory() as temp_dir:
            output_path = Path(temp_dir) / "exported_slm"
            # Export in native format
            export_result = await exporter.export(
                model_data=model_data,
                output_path=output_path,
                model_info=model_info,
                export_format="native",
            )
            # Verify export result
            assert export_result["format"] == "native"
            assert "files" in export_result
            assert len(export_result["files"]) == 5  # metadata.json, model.weights, etc.
            # Verify files were created
            assert (output_path / "metadata.json").exists()
            assert (output_path / "model.weights").exists()
            assert (output_path / "tokenizer.json").exists()
            # Step 3: Load the exported model
            runtime = SLMRuntime(config)
            loaded_model = await runtime.load_model(output_path)
            # Verify model was loaded correctly
            assert loaded_model.model_id == "gpt-3.5-turbo"
            assert loaded_model.parameters == 500000
            assert loaded_model.size == 250000
            # Step 4: Perform inference
            test_prompt = "Hello, how are you?"
            response = await runtime.generate(model=loaded_model, prompt=test_prompt, max_length=50)
            # Verify inference worked
            assert isinstance(response, str)
            assert len(response) > 0
            assert "[Generated by SLM" in response  # Check for SLM marker
            # Step 5: Benchmark the model
            benchmark_result = await runtime.benchmark(loaded_model)
            # Verify benchmark results
            assert "model_info" in benchmark_result
            assert "test_results" in benchmark_result
            assert "performance_metrics" in benchmark_result
            assert (
                benchmark_result["performance_metrics"]["total_tests"] == 5
            )  # Default test prompts

    @pytest.mark.asyncio
    async def test_multiple_format_export_load(self):
        """Test exporting and loading in different formats."""
        config = Config(openai_api_key="test_key")
        model_data = {
            "parameters": 100000,
            "size": 50000,
            "compression_ratio": 0.8,
        }
        model_info = {
            "model_id": "test-model",
            "provider": "test",
            "parameters": 1000000,
            "size": 500000,
        }
        exporter = SLMExporter(config)
        runtime = SLMRuntime(config)
        # Test native format
        with tempfile.TemporaryDirectory() as temp_dir:
            native_path = Path(temp_dir) / "native_model"
            await exporter.export(model_data, native_path, model_info, "native")
            native_model = await runtime.load_model(native_path)
            assert native_model.model_id == "test-model"
        # Test pickle format
        with tempfile.TemporaryDirectory() as temp_dir:
            pickle_dir = Path(temp_dir) / "pickle_model_dir"
            await exporter.export(model_data, pickle_dir, model_info, "pickle")
            pickle_file = pickle_dir / "model.pkl"
            pickle_model = await runtime.load_model(pickle_file)
            assert pickle_model.model_id == "test-model"

    @pytest.mark.asyncio
    async def test_model_caching_and_unloading(self):
        """Test model caching and unloading functionality."""
        config = Config(openai_api_key="test_key")
        model_data = {"parameters": 100000, "size": 50000}
        model_info = {"model_id": "cache-test", "provider": "test"}
        exporter = SLMExporter(config)
        runtime = SLMRuntime(config)
        with tempfile.TemporaryDirectory() as temp_dir:
            model_path = Path(temp_dir) / "cache_test"
            await exporter.export(model_data, model_path, model_info, "native")
            # Load model twice (should use cache)
            model1 = await runtime.load_model(model_path)
            model2 = await runtime.load_model(model_path)
            # Should be the same cached instance
            assert model1 is model2
            # Check loaded models list
            loaded_models = runtime.list_loaded_models()
            assert len(loaded_models) == 1
            assert loaded_models[0]["model_id"] == "cache-test"
            # Unload model
            result = runtime.unload_model(model_path)
            assert result is True
            # Should not be in loaded models anymore
            loaded_models = runtime.list_loaded_models()
            assert len(loaded_models) == 0

    @pytest.mark.asyncio
    async def test_inference_cache(self):
        """Test inference caching functionality."""
        config = Config(openai_api_key="test_key")
        model_data = {"parameters": 100000, "size": 50000}
        model_info = {"model_id": "inference-test", "provider": "test"}
        exporter = SLMExporter(config)
        runtime = SLMRuntime(config)
        with tempfile.TemporaryDirectory() as temp_dir:
            model_path = Path(temp_dir) / "inference_test"
            await exporter.export(model_data, model_path, model_info, "native")
            model = await runtime.load_model(model_path)
            prompt = "Test prompt"
            params = {"max_length": 50, "temperature": 0.7, "top_p": 0.9}
            # Generate response twice with same parameters
            response1 = await runtime.generate(model, prompt, **params)
            response2 = await runtime.generate(model, prompt, **params)
            # Should be the same (cached)
            assert response1 == response2
            # Clear cache and generate again
            runtime.clear_cache()
            response3 = await runtime.generate(model, prompt, **params)
            # Should be different (cache cleared)
            # Note: In real implementation, responses might still be similar due to
            # deterministic generation
            assert isinstance(response3, str)
            assert isinstance(response3, str)
