import torch
from torch import nn

from colipri.model.text import TextEncoder
from colipri.pooling import AttentionPool1D


class _StubBackbone(nn.Module):
    """Encode token IDs as deterministic embeddings."""

    def __init__(self):
        """Initialize embeddings with a distinctive padding-token vector."""
        super().__init__()
        weights = torch.tensor(
            [
                [100.0, -100.0, 50.0, -50.0],
                [1.0, 0.0, 0.0, 0.0],
                [0.0, 1.0, 0.0, 0.0],
                [0.0, 0.0, 1.0, 0.0],
            ]
        )
        self.embedding = nn.Embedding.from_pretrained(weights, freeze=False)

    def forward(self, token_ids, attention_mask=None):  # noqa: ARG002
        """Encode token IDs without altering padded positions."""
        return {"last_hidden_state": self.embedding(token_ids)}


def test_text_encoder_ignores_masked_padding():
    """Preserve the pooled embedding when masked padding is appended."""
    torch.manual_seed(0)
    encoder = TextEncoder(
        backbone=_StubBackbone(),
        pooler=AttentionPool1D(embed_dim=4, num_heads=2),
    )
    valid_ids = torch.tensor([[1, 2, 3]])
    padded_ids = torch.tensor([[1, 2, 3, 0, 0]])
    valid_mask = torch.ones_like(valid_ids)
    padded_mask = torch.tensor([[1, 1, 1, 0, 0]])

    expected = encoder(
        valid_ids,
        valid_mask,
        normalize=False,
    )
    result = encoder(
        padded_ids,
        padded_mask,
        normalize=False,
    )

    torch.testing.assert_close(result, expected)