"""Unit tests for kernel fusion module (Task3.4)."""

from __future__ import annotations

import numpy as np
import pytest

from sagellm_compression.fusion import (
    AttentionFusion,
    AttentionFusionConfig,
    FusionConfig,
    FusionType,
    GatedMLPFusionConfig,
    MLPFusion,
    MLPFusionConfig,
    StubKernelFusion,
)


def _rng() -> np.random.Generator:
    return np.random.default_rng(42)


def test_attention_fusion_end_to_end_with_packed_qkv() -> None:
    rng = _rng()
    cfg = AttentionFusionConfig(num_heads=4, head_dim=8, causal=True)
    fusion = AttentionFusion(cfg)

    batch, seq, hidden = 2, 16, cfg.hidden_dim
    hidden_states = rng.normal(size=(batch, seq, hidden)).astype(np.float32)
    qkv_weight = rng.normal(size=(3 * hidden, hidden)).astype(np.float32)
    out_weight = rng.normal(size=(hidden, hidden)).astype(np.float32)

    result = fusion.fuse(
        {
            "hidden_states": hidden_states,
            "qkv_weight": qkv_weight,
            "out_weight": out_weight,
        }
    )

    assert result.fusion_type == FusionType.ATTENTION
    assert result.device == "cpu"
    assert result.estimated_speedup >= 1.3
    assert result.is_accelerated
    assert result.ops_fused >= 2
    assert len(result.output["output"]) == batch * seq * hidden


def test_attention_fusion_with_separate_qkv_weights() -> None:
    rng = _rng()
    cfg = AttentionFusionConfig(num_heads=2, head_dim=8, causal=False)
    fusion = AttentionFusion(cfg)

    batch, seq, hidden = 1, 8, cfg.hidden_dim
    hidden_states = rng.normal(size=(batch, seq, hidden)).astype(np.float32)

    result = fusion.fuse(
        {
            "hidden_states": hidden_states,
            "q_weight": rng.normal(size=(hidden, hidden)).astype(np.float32),
            "k_weight": rng.normal(size=(hidden, hidden)).astype(np.float32),
            "v_weight": rng.normal(size=(hidden, hidden)).astype(np.float32),
            "out_weight": rng.normal(size=(hidden, hidden)).astype(np.float32),
        }
    )

    assert result.fusion_type == FusionType.ATTENTION
    assert len(result.output["output"]) == batch * seq * hidden


def test_attention_fusion_is_supported_false_for_bad_shape() -> None:
    cfg = AttentionFusionConfig(num_heads=2, head_dim=8)
    fusion = AttentionFusion(cfg)
    bad_hidden = np.zeros((2, 4, 15), dtype=np.float32)

    assert not fusion.is_supported({"hidden_states": bad_hidden})


def test_attention_fusion_fail_fast_for_missing_weights() -> None:
    cfg = AttentionFusionConfig(num_heads=2, head_dim=8)
    fusion = AttentionFusion(cfg)
    hidden_states = np.zeros((1, 4, cfg.hidden_dim), dtype=np.float32)

    with pytest.raises(ValueError, match="Missing required input"):
        fusion.fuse({"hidden_states": hidden_states})


def test_attention_fusion_bypass_when_disabled() -> None:
    cfg = AttentionFusionConfig(num_heads=2, head_dim=8, fuse_output_projection=False)
    cfg.enabled = False
    fusion = AttentionFusion(cfg)

    hidden_states = np.arange(1 * 4 * cfg.hidden_dim, dtype=np.float32).reshape(
        1, 4, cfg.hidden_dim
    )
    result = fusion.fuse({"hidden_states": hidden_states})

    assert result.metadata["bypassed"] is True
    assert result.estimated_speedup == 1.0
    assert result.ops_fused == 0
    assert result.output["output"] == hidden_states.ravel().tolist()


def test_mlp_fusion_standard_path() -> None:
    rng = _rng()
    cfg = MLPFusionConfig(hidden_dim=16, intermediate_dim=32, activation="gelu")
    fusion = MLPFusion(cfg)

    hidden_states = rng.normal(size=(2, 6, 16)).astype(np.float32)
    w1 = rng.normal(size=(32, 16)).astype(np.float32)
    w2 = rng.normal(size=(16, 32)).astype(np.float32)

    result = fusion.fuse({"hidden_states": hidden_states, "w1": w1, "w2": w2})

    assert result.fusion_type == FusionType.MLP
    assert result.estimated_speedup >= 1.3
    assert result.ops_fused == 3
    assert len(result.output["output"]) == 2 * 6 * 16


def test_mlp_fusion_gated_path() -> None:
    rng = _rng()
    cfg = GatedMLPFusionConfig(hidden_dim=8, intermediate_dim=12)
    fusion = MLPFusion(cfg)

    hidden_states = rng.normal(size=(1, 5, 8)).astype(np.float32)
    w_up = rng.normal(size=(12, 8)).astype(np.float32)
    w_down = rng.normal(size=(8, 12)).astype(np.float32)
    w_gate = rng.normal(size=(12, 8)).astype(np.float32)

    result = fusion.fuse(
        {
            "hidden_states": hidden_states,
            "w1": w_up,
            "w2": w_down,
            "w_gate": w_gate,
        }
    )

    assert result.fusion_type == FusionType.GATED_MLP
    assert result.metadata["gated"] is True
    assert result.ops_fused == 4


def test_mlp_fusion_fail_fast_for_missing_w1() -> None:
    cfg = MLPFusionConfig(hidden_dim=8, intermediate_dim=12)
    fusion = MLPFusion(cfg)
    hidden_states = np.zeros((1, 2, 8), dtype=np.float32)

    with pytest.raises(ValueError, match="Missing required input 'w1'"):
        fusion.fuse({"hidden_states": hidden_states, "w2": np.zeros((8, 12), dtype=np.float32)})


def test_mlp_fusion_fail_fast_for_hidden_dim_mismatch() -> None:
    cfg = MLPFusionConfig(hidden_dim=8, intermediate_dim=12)
    fusion = MLPFusion(cfg)

    with pytest.raises(ValueError, match="hidden_dim mismatch"):
        fusion.fuse(
            {
                "hidden_states": np.zeros((1, 2, 7), dtype=np.float32),
                "w1": np.zeros((12, 8), dtype=np.float32),
                "w2": np.zeros((8, 12), dtype=np.float32),
            }
        )


def test_stub_kernel_fusion_always_supported_and_accelerated() -> None:
    cfg = FusionConfig(fusion_type=FusionType.CUSTOM)
    fusion = StubKernelFusion(cfg)

    assert fusion.is_supported({})

    result = fusion.fuse({"x": [1.0, 2.0, 3.0]})
    assert result.estimated_speedup == 1.3
    assert result.is_accelerated
    assert result.metadata["stub"] is True


def test_fusion_config_validation_fail_fast() -> None:
    with pytest.raises(ValueError, match="device must be a non-empty string"):
        FusionConfig(device="")


def test_attention_config_validation_fail_fast() -> None:
    with pytest.raises(ValueError, match="num_heads must be >= 1"):
        AttentionFusionConfig(num_heads=0, head_dim=8)

    with pytest.raises(ValueError, match="head_dim must be >= 1"):
        AttentionFusionConfig(num_heads=4, head_dim=0)


def test_mlp_config_validation_fail_fast() -> None:
    with pytest.raises(ValueError, match="activation must be one of"):
        MLPFusionConfig(activation="not-real")  # type: ignore[arg-type]
