"""Lightweight text processor for QaDiT (T5 tokenizer wrapper)."""

from __future__ import annotations

from typing import List, Optional, Union

from transformers import AutoTokenizer
from transformers.feature_extraction_utils import BatchFeature

try:
    from .configuration_qadit import QaDiTConfig
except ImportError:
    from configuration_qadit import QaDiTConfig


class QaDiTProcessor:
    """Tokenize captions the same way training / ``generate`` expect.

    Usage::

        processor = QaDiTProcessor.from_pretrained("USER/qadit")
        inputs = processor("A dog barks", return_tensors="pt")
    """

    def __init__(self, tokenizer, text_max_length: int = 64):
        self.tokenizer = tokenizer
        self.text_max_length = text_max_length

    @classmethod
    def from_pretrained(cls, pretrained_model_name_or_path, **kwargs):
        trust = kwargs.pop("trust_remote_code", True)
        try:
            config = QaDiTConfig.from_pretrained(
                pretrained_model_name_or_path,
                trust_remote_code=trust,
            )
            text_model = config.text_model
            text_max_length = config.text_max_length
        except Exception:
            text_model = "google/flan-t5-large"
            text_max_length = 64
        tokenizer = AutoTokenizer.from_pretrained(text_model)
        return cls(tokenizer=tokenizer, text_max_length=text_max_length)

    def __call__(
        self,
        text: Union[str, List[str]],
        padding: str = "max_length",
        truncation: bool = True,
        max_length: Optional[int] = None,
        return_tensors: Optional[str] = "pt",
        **kwargs,
    ) -> BatchFeature:
        if isinstance(text, str):
            text = [text]
        encoded = self.tokenizer(
            text,
            padding=padding,
            truncation=truncation,
            max_length=max_length or self.text_max_length,
            return_tensors=return_tensors,
            **kwargs,
        )
        return BatchFeature(data=dict(encoded))

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

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


__all__ = ["QaDiTProcessor"]