from dataclasses import dataclass
from typing import Union
from PIL.Image import Image
from vllm import LLM, SamplingParams, TextPrompt

from src.modeling.llava.utils import load_tokenizer_and_processor
from src.modeling.multimodal import ImageTextGenerator, ImageTextModelConfig


# TODO: Remove the class and switch to ImageTextModelConfig directly.
@dataclass
class LlavaNextModelConfig(ImageTextModelConfig): ...


class LlavaNextGenerator(ImageTextGenerator):
    def __init__(self, config: Union[LlavaNextModelConfig, ImageTextModelConfig]) -> None:
        self.config = config
        self.tokenizer, self.processor = load_tokenizer_and_processor(
            config.model_name_or_path, use_fast=True
        )
        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,
        )

        super().__init__(model)

    def build_prompt(self, sample: str, images: Union[Image, list[Image], None]) -> str:
        num_images = (
            1 if isinstance(images, Image) else len(images) if isinstance(images, list) else 0
        )

        has_chat_template = hasattr(self.processor, "chat_template") and self.processor.chat_template is not None
        if not has_chat_template:
            # return "<image>\n" * num_images + sample
            return sample

        try:
            messages = [
                {
                    "role": "user",
                    "content": [
                        {"type": "text", "text": sample},
                        *[{"type": "image"}] * num_images,
                    ],
                }
            ]
            prompt = self.processor.apply_chat_template(
                messages, tokenize=False, add_generation_prompt=True
            )
        except:
            messages = [{"role": "user", "content": "<image>\n" * num_images + sample}]
            prompt = self.processor.apply_chat_template(
                messages, tokenize=False, add_generation_prompt=True
            )

        return prompt

    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.")

        prompts = []
        for sample, image in zip(samples, images):
            formatted_sample = self.config.prompt_template.format(text=sample)
            prompt = self.build_prompt(formatted_sample, image)
            prompts.append(
                TextPrompt(prompt=prompt, multi_modal_data={"image": image} if image else {})
            )

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

        out = self.model.generate(prompts, sampling_params=sampling_params)
        return [output.outputs[0].text.strip() for output in out]
