from typing import Any, Dict, Optional
import PIL
import torch
import PIL
import torch
from typing import Dict
from io import BytesIO
from transformers import SiglipImageProcessor
from sentence_transformers.models import Transformer as BaseTransformer


class MultiModalTransformer(BaseTransformer):

    def __init__(
        self,
        model_name_or_path: str,
        cache_dir: Optional[str] = None,
        tokenizer_args: Optional[Dict[str, Any]] = None,
        **kwargs,
    ):
        super().__init__(model_name_or_path, **kwargs)
        if tokenizer_args is None:
            tokenizer_args = {}
        self.processor = SiglipImageProcessor.from_pretrained(
            model_name_or_path, cache_dir=cache_dir, **tokenizer_args
        )

    def forward(
        self, features: dict[str, torch.Tensor], **kwargs
    ) -> dict[str, torch.Tensor]:
        trans_features = {
            "input_ids": features["input_ids"],
            "attention_mask": features["attention_mask"],
        }
        if "pixel_values" in features:
            trans_features["pixel_values"] = features["pixel_values"].to(
                self.auto_model.dtype
            )

        sentence_embedding = self.auto_model(**trans_features, **kwargs)[
            "sentence_embedding"
        ]
        features.update({"sentence_embedding": sentence_embedding})
        return features

    def tokenize(self, texts: list[Dict] | list[str]) -> dict[str, torch.Tensor]:
        img_start_token = "<|jasper_img_start|>"
        img_token = "<|jasper_img_token|>"
        img_end_token = "<|jasper_img_end|>"
        num_img_tokens = 300

        def process_text_item(item):
            if isinstance(item, str):
                return item, []
            text, images = "", []
            for sub_item in item:
                if sub_item["type"] == "text":
                    text += sub_item["content"]
                elif sub_item["type"] == "image_bytes":
                    text += img_start_token + img_token * num_img_tokens + img_end_token
                    images.append(
                        PIL.Image.open(BytesIO(sub_item["content"])).convert("RGB")
                    )
                elif sub_item["type"] == "image_path":
                    text += img_start_token + img_token * num_img_tokens + img_end_token
                    images.append(PIL.Image.open(sub_item["content"]).convert("RGB"))
                else:
                    raise ValueError(f"unknown data type {sub_item['type']}")
            return text, images

        all_texts, all_images = [], []
        for item in texts:
            text, images = process_text_item(item)
            all_texts.append(text)
            all_images.extend(images)
        ipt = self.tokenizer(
            all_texts,
            padding="longest",
            truncation=True,
            max_length=self.max_seq_length,
            return_tensors="pt",
        )
        if all_images:
            ipt["pixel_values"] = self.processor(
                images=all_images, return_tensors="pt"
            )["pixel_values"]
        return ipt