from __future__ import annotations

from typing import Any

import torch
from transformers import AutoConfig
from transformers.feature_extraction_utils import BatchFeature
from transformers.image_processing_utils import BaseImageProcessor
from transformers.processing_utils import ProcessorMixin
from transformers.tokenization_utils_base import PreTrainedTokenizerBase


class CapriProcessor(ProcessorMixin):
    attributes = ["image_processor", "tokenizer"]
    image_processor_class = "SiglipImageProcessor"
    tokenizer_class = ("Qwen2Tokenizer", "Qwen2TokenizerFast")

    def __init__(
        self,
        image_processor: BaseImageProcessor,
        tokenizer: PreTrainedTokenizerBase,
        prompt_prefix: str = "<image> Caption:",
        image_token: str = "<image>",
        pooled_embedding_dim: int = 768,
    ):
        self.prompt_prefix = prompt_prefix
        self.image_token = image_token
        self.pooled_embedding_dim = pooled_embedding_dim
        super().__init__(image_processor, tokenizer)

    @classmethod
    def from_pretrained(cls, pretrained_model_name_or_path, **kwargs):
        kwargs.setdefault("use_fast", False)
        processor = super().from_pretrained(pretrained_model_name_or_path, **kwargs)
        config = AutoConfig.from_pretrained(
            pretrained_model_name_or_path,
            trust_remote_code=kwargs.get("trust_remote_code", False),
        )
        processor.prompt_prefix = getattr(config, "prompt_prefix", processor.prompt_prefix)
        processor.image_token = getattr(config, "image_token", processor.image_token)
        processor.pooled_embedding_dim = getattr(config, "projector_in_dim", processor.pooled_embedding_dim)
        return processor

    def normalize_images(self, images) -> list[Any]:
        if isinstance(images, torch.Tensor):
            if images.ndim == 4:
                return [images[i] for i in range(images.shape[0])]
            return [images]
        if isinstance(images, (list, tuple)):
            return list(images)
        return [images]

    def normalize_pooled_embeddings(self, pooled_embeddings) -> torch.Tensor:
        pooled = torch.as_tensor(pooled_embeddings)
        if pooled.ndim == 1:
            pooled = pooled.unsqueeze(0)
        if pooled.ndim != 2:
            raise ValueError("`pooled_embeddings` must be a 1D embedding or a 2D batch of embeddings.")
        if pooled.shape[-1] != self.pooled_embedding_dim:
            raise ValueError(
                f"Expected pooled embedding dim {self.pooled_embedding_dim}, got {pooled.shape[-1]}."
            )
        return pooled

    def __call__(
        self,
        images=None,
        pooled_embeddings=None,
        text=None,
        return_tensors: str | None = "pt",
        padding: bool | str = True,
        truncation: bool = False,
        max_length: int | None = None,
        **kwargs: Any,
    ) -> BatchFeature:
        if images is None and pooled_embeddings is None and text is None:
            raise ValueError("Provide `images`, `pooled_embeddings`, or `text`.")

        batch = {}
        batch_size = None

        if images is not None:
            image_features = self.image_processor(images=images, return_tensors=return_tensors, **kwargs)
            batch.update(dict(image_features))
            batch_size = batch["pixel_values"].shape[0]

        if pooled_embeddings is not None:
            pooled = self.normalize_pooled_embeddings(pooled_embeddings)
            batch["pooled_embeddings"] = pooled
            batch_size = pooled.shape[0]

        if text is None and batch_size is not None:
            text = [self.prompt_prefix] * batch_size

        if text is not None:
            if isinstance(text, str):
                text = [text]
            tokenized = self.tokenizer(
                text,
                add_special_tokens=False,
                padding=padding,
                truncation=truncation,
                max_length=max_length,
                return_tensors=return_tensors,
            )
            batch.update(dict(tokenized))

        return BatchFeature(data=batch, tensor_type=return_tensors)

    def batch_decode(self, *args, **kwargs):
        return self.tokenizer.batch_decode(*args, **kwargs)

    def decode(self, *args, **kwargs):
        return self.tokenizer.decode(*args, **kwargs)