"""FineViT image processor wrapper."""

from __future__ import annotations

from transformers import AutoImageProcessor
from transformers.image_processing_utils import BaseImageProcessor


class FineViTImageProcessor(BaseImageProcessor):
    model_input_names = ["pixel_values"]

    def __init__(
        self,
        backbone_model_name: str = "facebook/dinov2-with-registers-base",
        image_size: int = 224,
        **kwargs,
    ):
        super().__init__(**kwargs)
        self.backbone_model_name = backbone_model_name
        self.image_size = int(image_size)
        self._backbone_processor = None

    @property
    def backbone_processor(self):
        if self._backbone_processor is None:
            processor = AutoImageProcessor.from_pretrained(self.backbone_model_name)
            self._set_square_size(processor)
            self._backbone_processor = processor
        return self._backbone_processor

    @property
    def image_mean(self):
        return getattr(self.backbone_processor, "image_mean", [0.485, 0.456, 0.406])

    @property
    def image_std(self):
        return getattr(self.backbone_processor, "image_std", [0.229, 0.224, 0.225])

    def _set_square_size(self, processor) -> None:
        size = {"height": self.image_size, "width": self.image_size}
        if hasattr(processor, "size"):
            current = getattr(processor, "size")
            if isinstance(current, dict) and "shortest_edge" in current:
                processor.size = {"shortest_edge": self.image_size}
            else:
                processor.size = size
        if hasattr(processor, "crop_size"):
            processor.crop_size = size

    def __call__(self, images, **kwargs):
        return self.backbone_processor(images=images, **kwargs)

    def to_dict(self):
        output = super().to_dict()
        output.pop("_backbone_processor", None)
        return output


FineViTImageProcessor.register_for_auto_class("AutoImageProcessor")