"""Processor for CAIP: bundles the image processor and the SigLIP 2 tokenizer.

Mirrors open_clip's `HFTokenizer.__call__` for the text path: each string is
`canonicalize`-cleaned (underscores->spaces, punctuation removed, lowercased,
whitespace collapsed) then tokenized with padding to a fixed context length
(64) and truncation. The pad token id (0) is what the model's cross-attention
pool treats as padding.

    proc = AutoProcessor.from_pretrained(repo, trust_remote_code=True)
    inputs = proc(images=pil_image, text="pick up the red cup", return_tensors="pt")
    # inputs["pixel_values"] [B,3,256,256], inputs["input_ids"] [B,64]
"""

import string

from transformers.feature_extraction_utils import BatchFeature
from transformers.processing_utils import ProcessorMixin

_PUNCT_TABLE = str.maketrans("", "", string.punctuation)


def canonicalize_text(text: str) -> str:
    # From open_clip.tokenizer.canonicalize_text (big_vision canonicalization).
    text = text.replace("_", " ")
    text = text.translate(_PUNCT_TABLE)
    text = text.lower()
    text = " ".join(text.split())
    return text.strip()


class CaipProcessor(ProcessorMixin):
    attributes = ["image_processor", "tokenizer"]
    image_processor_class = "AutoImageProcessor"
    tokenizer_class = "AutoTokenizer"

    def __init__(self, image_processor=None, tokenizer=None, context_length: int = 64, **kwargs):
        super().__init__(image_processor, tokenizer)
        self.context_length = context_length

    def __call__(self, images=None, text=None, return_tensors="pt", **kwargs):
        if images is None and text is None:
            raise ValueError("Provide at least one of `images` or `text`.")
        data = {}
        if images is not None:
            data.update(self.image_processor(images, return_tensors=return_tensors))
        if text is not None:
            if isinstance(text, str):
                text = [text]
            text = [canonicalize_text(t) for t in text]
            enc = self.tokenizer(
                text,
                return_tensors=return_tensors,
                max_length=self.context_length,
                padding="max_length",
                truncation=True,
                add_special_tokens=True,
            )
            data["input_ids"] = enc["input_ids"]
        return BatchFeature(data=data, tensor_type=return_tensors)