import pytest
import torch

from mini_transformer.modules.decoder import DecoderLayer, TransformerDecoder


def _dev():
    return (
        torch.device(f"cuda:{torch.cuda.current_device()}")
        if torch.cuda.is_available()
        else torch.device("cpu")
    )


# -------- ctor validations --------


def test_decoder_layer_param_checks():
    with pytest.raises(TypeError):
        DecoderLayer("32", 4, 64, 0.1)
    with pytest.raises(TypeError):
        DecoderLayer(32, "4", 64, 0.1)
    with pytest.raises(TypeError):
        DecoderLayer(32, 4, "64", 0.1)
    with pytest.raises(TypeError):
        DecoderLayer(32, 4, 64, "0.1")
    with pytest.raises(TypeError):
        DecoderLayer(32, 4, 64, 0.1, layer_norm_style=123)

    with pytest.raises(ValueError):
        DecoderLayer(0, 4, 64, 0.1)
    with pytest.raises(ValueError):
        DecoderLayer(32, 0, 64, 0.1)
    with pytest.raises(ValueError):
        DecoderLayer(32, 4, 0, 0.1)
    with pytest.raises(ValueError):
        DecoderLayer(32, 4, 64, 1.0)  # upper bound excluded
    with pytest.raises(ValueError):
        DecoderLayer(32, 4, 64, 0.1, layer_norm_style="weird")


def test_transformer_decoder_layers_count_checks():
    with pytest.raises(TypeError):
        TransformerDecoder(32, 4, 64, "2", 0.1)
    with pytest.raises(ValueError):
        TransformerDecoder(32, 4, 64, 0, 0.1)


# -------- forward path --------


@pytest.mark.parametrize("B,Sx,Sy,D,H,FF,L", [(2, 5, 6, 24, 3, 48, 2)])
def test_decoder_forward_happy_path(B, Sx, Sy, D, H, FF, L):
    device = _dev()
    dec = TransformerDecoder(D, H, FF, L, 0.1).to(device)

    x = torch.randn(B, Sx, D, device=device)
    y = torch.randn(B, Sy, D, device=device)
    heads = dec.layers[0].self_attention_layer.num_heads
    src_pad = torch.zeros(B, heads, 1, Sx, dtype=torch.bool, device=device)
    tgt_pad = torch.zeros(B, heads, 1, Sy, dtype=torch.bool, device=device)
    causal = (
        torch.ones(B, heads, Sy, Sy, dtype=torch.bool, device=device).triu(1) if Sy > 0 else None
    )

    out = dec(x, y, src_pad, tgt_pad, causal)
    assert out.shape == (B, Sy, D)
    assert out.device == device


def test_transformer_decoder_pre_norm_layers_flag():
    dec = TransformerDecoder(24, 3, 48, 2, 0.1, layer_norm_style="pre")
    assert all(layer.pre_norm for layer in dec.layers)


def test_decoder_layer_pre_norm_forward_shapes():
    device = _dev()
    layer = DecoderLayer(24, 3, 48, 0.1, layer_norm_style="pre").to(device)
    x = torch.randn(2, 5, 24, device=device)
    y = torch.randn(2, 6, 24, device=device)
    heads = layer.self_attention_layer.num_heads
    src_mask = torch.zeros(2, heads, 1, 5, dtype=torch.bool, device=device)
    tgt_mask = torch.zeros(2, heads, 1, 6, dtype=torch.bool, device=device)
    causal = torch.ones(2, heads, 6, 6, dtype=torch.bool, device=device).triu(1)
    out = layer(x, y, src_mask, tgt_mask, causal)
    assert out.shape == y.shape
    assert layer.pre_norm is True


def test_decoder_forward_input_checks_and_message_format():
    layer = DecoderLayer(24, 3, 48, 0.1)
    heads = layer.self_attention_layer.num_heads
    base_src_mask = torch.zeros(2, heads, 1, 3, dtype=torch.bool)
    base_tgt_mask = torch.zeros(2, heads, 1, 3, dtype=torch.bool)

    with pytest.raises(TypeError):
        layer("not a tensor", torch.randn(2, 3, 24), base_src_mask, base_tgt_mask, None)
    with pytest.raises(TypeError):
        layer(torch.randn(2, 3, 24), "not a tensor", base_src_mask, base_tgt_mask, None)

    with pytest.raises(ValueError) as e1:
        layer(
            torch.randn(2, 3, 24, 5),
            torch.randn(2, 3, 24),
            base_src_mask,
            base_tgt_mask,
            None,
        )  # x rank 4
    assert "x must be a 3D torch.Tensor of shape (B, S, D)" in str(e1.value)

    with pytest.raises(ValueError) as e2:
        layer(
            torch.randn(2, 3, 24),
            torch.randn(2, 3, 24, 5),
            base_src_mask,
            base_tgt_mask,
            None,
        )  # y rank 4
    assert "y must be a 3D torch.Tensor of shape (B, S, D)" in str(e2.value)

    x = torch.randn(2, 5, 24)
    y = torch.randn(3, 6, 24)  # batch mismatch
    src_mask_x = torch.zeros(2, heads, 1, 5, dtype=torch.bool)
    tgt_mask_y = torch.zeros(3, heads, 1, 6, dtype=torch.bool)
    with pytest.raises(ValueError) as e3:
        layer(x, y, src_mask_x, tgt_mask_y, None)
    assert "Encoder memory and decoder input must match in batch and d_model" in str(e3.value)

    y2 = torch.randn(2, 6, 16)  # d_model mismatch
    tgt_mask_y2 = torch.zeros(2, heads, 1, 6, dtype=torch.bool)
    with pytest.raises(ValueError) as e4:
        layer(x, y2, src_mask_x, tgt_mask_y2, None)
    assert "Encoder memory and decoder input must match in batch and d_model" in str(e4.value)