"""Tests for compression_types module (issue #20 / #21)"""

from __future__ import annotations

import pytest
from pydantic import ValidationError

from sagellm_protocol import (
    CompressionConfig,
    CompressionMetrics,
    FusionConfig,
    FusionLevel,
    QuantizationConfig,
    QuantizationGranularity,
    QuantizationScheme,
    SparseConfig,
    SparseFormat,
    SparsityPattern,
    SpeculativeDecodingConfig,
    SpeculativeStrategy,
)


class TestQuantizationScheme:
    def test_all_schemes_exist(self) -> None:
        assert QuantizationScheme.NONE == "none"
        assert QuantizationScheme.INT8 == "int8"
        assert QuantizationScheme.FP8 == "fp8"
        assert QuantizationScheme.INT4 == "int4"
        assert QuantizationScheme.NF4 == "nf4"
        assert QuantizationScheme.AWQ == "awq"
        assert QuantizationScheme.GPTQ == "gptq"
        assert QuantizationScheme.SMOOTH_QUANT == "smooth_quant"
        assert QuantizationScheme.BITSANDBYTES == "bitsandbytes"
        assert QuantizationScheme.GGUF == "gguf"


class TestQuantizationConfig:
    def test_defaults(self) -> None:
        cfg = QuantizationConfig()
        assert cfg.scheme == QuantizationScheme.NONE
        assert cfg.weight_granularity == QuantizationGranularity.PER_CHANNEL
        assert cfg.activation_granularity == QuantizationGranularity.PER_TENSOR
        assert cfg.group_size is None
        assert cfg.bits is None

    def test_int8_config(self) -> None:
        cfg = QuantizationConfig(
            scheme=QuantizationScheme.INT8,
            bits=8,
            weight_granularity=QuantizationGranularity.PER_CHANNEL,
        )
        assert cfg.scheme == QuantizationScheme.INT8
        assert cfg.bits == 8

    def test_awq_group_quant(self) -> None:
        cfg = QuantizationConfig(
            scheme=QuantizationScheme.AWQ,
            bits=4,
            weight_granularity=QuantizationGranularity.PER_GROUP,
            group_size=128,
        )
        assert cfg.group_size == 128

    def test_fp8_with_format(self) -> None:
        cfg = QuantizationConfig(
            scheme=QuantizationScheme.FP8,
            fp8_format="e4m3",
        )
        assert cfg.fp8_format == "e4m3"

    def test_bits_range(self) -> None:
        # valid range: 2..32
        QuantizationConfig(bits=2)
        QuantizationConfig(bits=32)
        with pytest.raises(ValidationError):
            QuantizationConfig(bits=1)
        with pytest.raises(ValidationError):
            QuantizationConfig(bits=33)

    def test_group_size_must_be_positive(self) -> None:
        with pytest.raises(ValidationError):
            QuantizationConfig(group_size=0)


class TestSparseConfig:
    def test_defaults(self) -> None:
        cfg = SparseConfig()
        assert cfg.enabled is False
        assert cfg.pattern == SparsityPattern.UNSTRUCTURED
        assert cfg.sparse_format == SparseFormat.DENSE
        assert cfg.sparsity_ratio == 0.0

    def test_n_m_sparsity(self) -> None:
        cfg = SparseConfig(
            enabled=True,
            pattern=SparsityPattern.SEMI_STRUCTURED,
            sparse_format=SparseFormat.N_M,
            n=2,
            m=4,
            sparsity_ratio=0.5,
        )
        assert cfg.n == 2
        assert cfg.m == 4
        assert cfg.sparsity_ratio == 0.5

    def test_sparsity_ratio_range(self) -> None:
        SparseConfig(sparsity_ratio=0.0)
        SparseConfig(sparsity_ratio=1.0)
        with pytest.raises(ValidationError):
            SparseConfig(sparsity_ratio=-0.1)
        with pytest.raises(ValidationError):
            SparseConfig(sparsity_ratio=1.1)

    def test_block_sparse(self) -> None:
        cfg = SparseConfig(
            enabled=True,
            pattern=SparsityPattern.BLOCK,
            sparse_format=SparseFormat.BSR,
            block_size=16,
        )
        assert cfg.block_size == 16

    def test_all_sparse_formats_exist(self) -> None:
        for fmt in SparseFormat:
            SparseConfig(sparse_format=fmt)


class TestSpeculativeDecodingConfig:
    def test_defaults(self) -> None:
        cfg = SpeculativeDecodingConfig()
        assert cfg.enabled is False
        assert cfg.strategy == SpeculativeStrategy.DRAFT_MODEL
        assert cfg.num_speculative_tokens == 4
        assert cfg.acceptance_threshold == 0.0

    def test_draft_model_config(self) -> None:
        cfg = SpeculativeDecodingConfig(
            enabled=True,
            strategy=SpeculativeStrategy.DRAFT_MODEL,
            draft_model="meta-llama/Llama-3.2-1B",
            num_speculative_tokens=5,
        )
        assert cfg.draft_model == "meta-llama/Llama-3.2-1B"
        assert cfg.num_speculative_tokens == 5

    def test_medusa_config(self) -> None:
        cfg = SpeculativeDecodingConfig(
            enabled=True,
            strategy=SpeculativeStrategy.MEDUSA,
            num_heads=4,
            num_speculative_tokens=3,
        )
        assert cfg.num_heads == 4

    def test_num_speculative_tokens_range(self) -> None:
        SpeculativeDecodingConfig(num_speculative_tokens=1)
        SpeculativeDecodingConfig(num_speculative_tokens=64)
        with pytest.raises(ValidationError):
            SpeculativeDecodingConfig(num_speculative_tokens=0)
        with pytest.raises(ValidationError):
            SpeculativeDecodingConfig(num_speculative_tokens=65)

    def test_acceptance_threshold_range(self) -> None:
        SpeculativeDecodingConfig(acceptance_threshold=0.0)
        SpeculativeDecodingConfig(acceptance_threshold=1.0)
        with pytest.raises(ValidationError):
            SpeculativeDecodingConfig(acceptance_threshold=1.1)

    def test_all_strategies_exist(self) -> None:
        for s in SpeculativeStrategy:
            SpeculativeDecodingConfig(strategy=s)


class TestFusionConfig:
    def test_defaults(self) -> None:
        cfg = FusionConfig()
        assert cfg.level == FusionLevel.NONE
        assert cfg.use_flash_attention is False

    def test_full_fusion(self) -> None:
        cfg = FusionConfig(
            level=FusionLevel.FULL,
            use_flash_attention=True,
            use_fused_rope=True,
            use_fused_layernorm=True,
            use_fused_swiglu=True,
        )
        assert cfg.level == FusionLevel.FULL
        assert cfg.use_flash_attention is True
        assert cfg.use_fused_swiglu is True


class TestCompressionConfig:
    def test_defaults(self) -> None:
        cfg = CompressionConfig()
        assert cfg.quantization.scheme == QuantizationScheme.NONE
        assert cfg.sparsity.enabled is False
        assert cfg.speculative_decoding.enabled is False
        assert cfg.fusion.level == FusionLevel.NONE

    def test_full_compression_config(self) -> None:
        cfg = CompressionConfig(
            quantization=QuantizationConfig(scheme=QuantizationScheme.AWQ, bits=4),
            sparsity=SparseConfig(enabled=True, sparsity_ratio=0.5),
            speculative_decoding=SpeculativeDecodingConfig(
                enabled=True,
                strategy=SpeculativeStrategy.DRAFT_MODEL,
                draft_model="draft-model",
            ),
            fusion=FusionConfig(level=FusionLevel.ATTENTION, use_flash_attention=True),
        )
        assert cfg.quantization.scheme == QuantizationScheme.AWQ
        assert cfg.sparsity.enabled is True
        assert cfg.speculative_decoding.enabled is True
        assert cfg.fusion.level == FusionLevel.ATTENTION

    def test_serialization_roundtrip(self) -> None:
        cfg = CompressionConfig(
            quantization=QuantizationConfig(scheme=QuantizationScheme.GPTQ, bits=4, group_size=128),
        )
        data = cfg.model_dump()
        cfg2 = CompressionConfig.model_validate(data)
        assert cfg2.quantization.scheme == QuantizationScheme.GPTQ


class TestCompressionMetrics:
    def test_defaults(self) -> None:
        m = CompressionMetrics()
        assert m.quant_scheme == QuantizationScheme.NONE
        assert m.compress_ratio == 1.0
        assert m.quant_error_pct == 0.0
        assert m.actual_sparsity_ratio == 0.0
        assert m.fusion_level == FusionLevel.NONE

    def test_int8_metrics(self) -> None:
        m = CompressionMetrics(
            quant_scheme=QuantizationScheme.INT8,
            compress_ratio=2.0,
            quant_error_pct=0.05,
        )
        assert m.compress_ratio == 2.0
        assert m.quant_error_pct == 0.05

    def test_spec_decoding_metrics(self) -> None:
        m = CompressionMetrics(
            spec_tokens_proposed=100,
            spec_tokens_accepted=85,
        )
        assert m.spec_tokens_accepted == 85

    def test_compress_ratio_must_be_ge_1(self) -> None:
        with pytest.raises(ValidationError):
            CompressionMetrics(compress_ratio=0.5)

    def test_sparsity_ratio_range(self) -> None:
        with pytest.raises(ValidationError):
            CompressionMetrics(actual_sparsity_ratio=1.5)
