from typing import Union, List
from io import BytesIO
import base64
from PIL.Image import Image
from vllm import LLM, SamplingParams

from src.modeling.multimodal import ImageTextGenerator, ImageTextModelConfig


class MistralGenerator(ImageTextGenerator):
    def __init__(self, config: ImageTextModelConfig) -> None:
        self.config = config
        model = LLM(
            model=config.model_name_or_path,
            trust_remote_code=True,
            gpu_memory_utilization=config.gpu_memory_utilization,
            dtype="bfloat16",
            enforce_eager=config.enforce_eager,
            max_model_len=config.max_seq_len,
            tensor_parallel_size=config.num_gpus,
            tokenizer_mode="mistral",
        )
        super().__init__(model)

    def _build_conversation(
        self,
        text: str,
        images: Union[Image, List[Image], None],
    ) -> list[dict]:
        def image_to_url(image: Image) -> str:
            fmt = image.format if image.format else "PNG"
            buf = BytesIO()
            image.save(buf, format=fmt)
            b64 = base64.b64encode(buf.getvalue()).decode("utf-8")
            return f"data:image/{fmt.lower()};base64,{b64}"

        content = [{"type": "text", "text": text}]
        for image in images if isinstance(images, list) else [images]:
            if image is None:
                continue

            image_content = {"type": "image_url", "image_url": {"url": image_to_url(image)}}
            content.append(image_content)
        return [{"role": "user", "content": content}]

    def generate(
        self,
        samples: list[str],
        images: list[Union[Image, list[Image], None]]
    ) -> list[str]:
        if len(samples) != len(images):
            raise ValueError("Samples and images must have the same length.")

        sampling_params = SamplingParams(
            temperature=self.config.temperature,
            top_p=self.config.top_p,
            max_tokens=self.config.max_new_tokens,
        )

        conversations = []
        for sample, imgs in zip(samples, images):
            formatted = self.config.prompt_template.format(text=sample)
            conversations.append(self._build_conversation(formatted, imgs))

        results = self.model.chat(messages=conversations, sampling_params=sampling_params)
        outputs = [r.outputs[0].text.strip() for r in results]
        return outputs
