"""
tests/test_quantization_matrix.py — Gap #2: Quantization Robustness Testing.

Tests QuantizationTestMatrix, VarianceMetrics, and ModelAdapter.load_quantized().
All tests run without loading real models (callable adapters + mocking).
"""
import pytest
from unittest.mock import MagicMock, patch

from sichgate.testing.quantization_matrix import (
    QuantizationTestMatrix,
    QuantizationVarianceReport,
    AttackVarianceRow,
    VarianceMetrics,
    AttackQuantizationResult,
    QUANTIZATION_LEVELS,
    _HIGH_VARIANCE_THRESHOLD,
)
from sichgate.models.adapter import ModelAdapter


# ---------------------------------------------------------------------------
# Unit tests: VarianceMetrics / compute_variance
# ---------------------------------------------------------------------------

class TestComputeVariance:
    def test_high_variance_detection(self):
        matrix = QuantizationTestMatrix("test-model")
        metrics = matrix.compute_variance({"Q4": 0.70, "Q8": 0.44, "FP16": 0.42})

        assert metrics.max_delta == pytest.approx(0.28, abs=0.01)
        assert metrics.significant is True
        assert metrics.trend == "higher_risk_lower_precision"
        assert metrics.min_level == "FP16"
        assert metrics.max_level == "Q4"

    def test_low_variance_flat(self):
        matrix = QuantizationTestMatrix("test-model")
        metrics = matrix.compute_variance({"Q4": 0.50, "Q8": 0.51, "FP16": 0.52})

        assert metrics.max_delta < _HIGH_VARIANCE_THRESHOLD
        assert metrics.significant is False
        assert metrics.trend == "flat"

    def test_higher_risk_higher_precision(self):
        matrix = QuantizationTestMatrix("test-model")
        # Unusual case: FP16 more vulnerable than Q4
        metrics = matrix.compute_variance({"Q4": 0.20, "FP16": 0.65})

        assert metrics.trend == "higher_risk_higher_precision"

    def test_empty_input_returns_flat(self):
        matrix = QuantizationTestMatrix("test-model")
        metrics = matrix.compute_variance({})

        assert metrics.max_delta == 0.0
        assert metrics.significant is False
        assert metrics.trend == "flat"

    def test_single_level_no_stddev(self):
        matrix = QuantizationTestMatrix("test-model")
        metrics = matrix.compute_variance({"FP16": 0.45})

        assert metrics.max_delta == 0.0
        assert metrics.stddev == 0.0

    def test_variance_to_dict(self):
        matrix = QuantizationTestMatrix("test-model")
        metrics = matrix.compute_variance({"Q4": 0.80, "FP16": 0.20})
        d = metrics.to_dict()

        assert "max_delta" in d
        assert "stddev" in d
        assert "trend" in d
        assert "significant" in d


# ---------------------------------------------------------------------------
# Unit tests: recommend_quantization
# ---------------------------------------------------------------------------

class TestRecommendQuantization:
    def _make_row(self, name: str, max_level: str, significant: bool) -> AttackVarianceRow:
        return AttackVarianceRow(
            attack_name=name,
            attack_type="llm",
            by_quantization={},
            variance=VarianceMetrics(
                max_delta=0.30 if significant else 0.05,
                stddev=0.10,
                min_level="FP16",
                max_level=max_level,
                trend="higher_risk_lower_precision",
                significant=significant,
            ),
        )

    def test_no_variance_positive_recommendation(self):
        matrix = QuantizationTestMatrix("model")
        rows = [self._make_row("attack1", "FP16", False)]
        report = QuantizationVarianceReport(
            model="model", quantization_variants=["Q4", "FP16"],
            results=rows, recommendation="",
            created_at="", high_variance_attacks=[],
        )
        rec = matrix.recommend_quantization(report)
        assert "safe" in rec.lower() or "any" in rec.lower()

    def test_q4_high_risk_recommends_fp16(self):
        matrix = QuantizationTestMatrix("model")
        rows = [self._make_row("jailbreak", "Q4", True)]
        report = QuantizationVarianceReport(
            model="model", quantization_variants=["Q4", "FP16"],
            results=rows, recommendation="",
            created_at="", high_variance_attacks=["jailbreak"],
        )
        rec = matrix.recommend_quantization(report)
        assert "Q4" in rec or "avoid" in rec.lower()


# ---------------------------------------------------------------------------
# Unit tests: ModelAdapter.load_quantized()
# ---------------------------------------------------------------------------

class TestModelAdapterLoadQuantized:
    def test_load_quantized_returns_adapter(self):
        """load_quantized() returns a ModelAdapter with quantization_level set."""
        adapter = ModelAdapter.load_quantized("test-model", "FP16")

        assert isinstance(adapter, ModelAdapter)
        assert adapter.model_path == "test-model"
        assert adapter.adapter_type == "hf"
        assert getattr(adapter, "_quantization_level", None) == "FP16"

    def test_load_quantized_invalid_level_raises(self):
        """Unknown quantization levels raise ValueError."""
        with pytest.raises(ValueError, match="Q2"):
            ModelAdapter.load_quantized("test-model", "Q2")

    def test_load_quantized_passes_hf_token(self):
        adapter = ModelAdapter.load_quantized(
            "private/model", "Q8", hf_token="hf_test_token"
        )
        assert adapter._hf_token == "hf_test_token"

    def test_load_quantized_callable_override(self):
        """Quantized adapter can be wrapped as callable for testing without real model."""
        def mock_fn(prompt: str) -> str:
            return "test response"

        adapter = ModelAdapter(
            model_path="test-model",
            adapter_type="callable",
            callable_fn=mock_fn,
        )
        assert adapter.is_reachable()
        assert "test response" == adapter.query("hello")

    def test_all_quantization_levels_accepted(self):
        """All 4 defined levels are valid."""
        for level in ["Q4", "Q8", "INT8", "FP16"]:
            adapter = ModelAdapter.load_quantized("model", level)
            assert getattr(adapter, "_quantization_level") == level


# ---------------------------------------------------------------------------
# Integration tests: QuantizationTestMatrix with mocked runner
# ---------------------------------------------------------------------------

class TestQuantizationTestMatrixIntegration:
    def _mock_attack_cls(self, name: str, success_rate: float):
        cls = MagicMock()
        cls.__name__ = name
        cls.attack_name = name
        cls.attack_type = "llm"
        return cls

    def _make_mock_summary(self, attack_name: str, success_rate: float):
        result = {
            "attack_name": attack_name, "attack_type": "llm",
            "success_rate": success_rate, "status": "fail",
            "severity": "high",
        }
        mock_summary = MagicMock()
        mock_summary.results = [result]
        return mock_summary

    @patch("sichgate.testing.quantization_matrix.QuantizationTestMatrix._load_quantized_adapter")
    def test_execute_attack_matrix_returns_report(self, mock_load_adapter):
        """execute_attack_matrix() returns a QuantizationVarianceReport."""
        mock_adapter = MagicMock()
        mock_adapter.model_path = "test-model"
        mock_adapter.adapter_type = "callable"
        mock_load_adapter.return_value = mock_adapter

        # Inject mock runner module so sichgate.runner doesn't need to be imported
        call_counter = {"n": 0}
        mock_runner_instance = MagicMock()
        def mock_run(model_path, attack_classes, adapter_type):
            call_counter["n"] += 1
            rate = 0.70 if call_counter["n"] == 1 else 0.40
            return self._make_mock_summary("test_attack", rate)
        mock_runner_instance.run.side_effect = mock_run

        mock_runner_cls = MagicMock(return_value=mock_runner_instance)
        mock_runner_module = MagicMock()
        mock_runner_module.AttackRunner = mock_runner_cls

        attack_cls = self._mock_attack_cls("test_attack", 0.5)
        matrix = QuantizationTestMatrix("test-model")

        with patch.dict("sys.modules", {"sichgate.runner": mock_runner_module}):
            report = matrix.execute_attack_matrix([attack_cls], levels=["Q4", "FP16"])

        assert isinstance(report, QuantizationVarianceReport)
        assert set(report.quantization_variants).issubset({"Q4", "FP16"})
        assert len(report.results) == 1

    def test_invalid_level_raises_value_error(self):
        """execute_attack_matrix() raises ValueError for unknown levels."""
        matrix = QuantizationTestMatrix("test-model")
        with pytest.raises(ValueError, match="Q2"):
            matrix.execute_attack_matrix([], levels=["Q2"])

    @patch("sichgate.testing.quantization_matrix.QuantizationTestMatrix._load_quantized_adapter")
    def test_no_levels_load_fails_raises(self, mock_load_adapter):
        """If all adapter loads fail, RuntimeError is raised."""
        mock_load_adapter.return_value = None  # All levels fail to load

        matrix = QuantizationTestMatrix("test-model")
        with pytest.raises(RuntimeError, match="No quantization levels"):
            matrix.execute_attack_matrix([], levels=["FP16"])

    def test_report_to_dict_serializable(self):
        """QuantizationVarianceReport.to_dict() produces JSON-serializable output."""
        import json
        rows = [
            AttackVarianceRow(
                attack_name="jailbreak",
                attack_type="llm",
                by_quantization={
                    "Q4": AttackQuantizationResult("Q4", 0.7, 1.0, {"fail": 1}),
                    "FP16": AttackQuantizationResult("FP16", 0.4, 1.0, {"fail": 1}),
                },
                variance=VarianceMetrics(
                    max_delta=0.30, stddev=0.15,
                    min_level="FP16", max_level="Q4",
                    trend="higher_risk_lower_precision", significant=True,
                ),
            )
        ]
        report = QuantizationVarianceReport(
            model="tinyllama",
            quantization_variants=["Q4", "FP16"],
            results=rows,
            recommendation="Deploy FP16",
            created_at="2026-02-27T00:00:00Z",
            high_variance_attacks=["jailbreak"],
        )
        d = report.to_dict()
        serialized = json.dumps(d)  # Should not raise
        assert "tinyllama" in serialized
        assert "jailbreak" in serialized
