from collections import UserDict

import torch
from torch import nn

from cortex_ai.vision.native import build_native_fusion


class TinyCore(nn.Module):
    def __init__(self):
        super().__init__()
        self.embedding = nn.Embedding(32, 8)
        self.last_inputs_embeds = None

    def get_input_embeddings(self):
        return self.embedding

    def forward(self, inputs_embeds, attention_mask):
        self.last_inputs_embeds = inputs_embeds
        return {"hidden_states": inputs_embeds, "attention_mask": attention_mask}


class TinyVisionEncoder(nn.Module):
    def forward(self, image):
        return torch.ones(image.shape[0], 3, 4)


def test_native_image_embedding_fusion_executes_real_tensor_path():
    core = TinyCore()
    fusion = build_native_fusion(TinyVisionEncoder(), core, vision_hidden_size=4, cortex_hidden_size=8)
    result = fusion(input_ids=torch.tensor([[1, 2]]), image=torch.zeros(1, 3, 16, 16))
    assert result.image_token_count == 3
    assert result.text_token_count == 2
    assert core.last_inputs_embeds.shape == (1, 5, 8)
    assert result.outputs["attention_mask"].shape == (1, 5)


class BatchFeatureLikeEncoder(nn.Module):
    """Encoder whose processor output is a UserDict, as Transformers processors return."""

    def forward(self, pixel_values):
        return torch.ones(pixel_values.shape[0], 2, 4)


class BatchFeatureLikeProcessor:
    def __call__(self, image, return_tensors=None):
        return UserDict({"pixel_values": torch.zeros(1, 3, 16, 16)})


def test_processor_output_is_unpacked_by_keyword_not_positionally():
    core = TinyCore()
    fusion = build_native_fusion(
        BatchFeatureLikeEncoder(),
        core,
        vision_hidden_size=4,
        cortex_hidden_size=8,
        image_processor=BatchFeatureLikeProcessor(),
    )
    result = fusion(input_ids=torch.tensor([[1, 2]]), image="unused")
    assert result.image_token_count == 2
    assert result.text_token_count == 2