from concurrent.futures import ThreadPoolExecutor
from typing import List, Union

import torch
from PIL.Image import Image
from tqdm import tqdm
from transformers import AutoProcessor, PaliGemmaForConditionalGeneration
from vllm import LLM, SamplingParams

from src.logger import get_logger
from src.utils.iterators import batch_with_preload
from src.modeling.multimodal import ImageTextGenerator, ImageTextModelConfig


logger = get_logger(__name__)


class PaligemmaGenerator(ImageTextGenerator):
    def __init__(self, config: ImageTextModelConfig) -> None:
        self.config = config
        self.device = "cuda" if torch.cuda.is_available() else "cpu"

        # Load Processor
        self.processor = AutoProcessor.from_pretrained(
            config.model_name_or_path,
            trust_remote_code=True,
        )

        # Load Model
        self.model = PaliGemmaForConditionalGeneration.from_pretrained(
            config.model_name_or_path,
            trust_remote_code=True,
            torch_dtype=torch.bfloat16,
            device_map=self.device,
        ).eval()

        super().__init__(self.model)

    def generate(
        self, samples: List[str], images: List[Union[Image, List[Image], None]]
    ) -> List[str]:
        outputs: List[str] = []
        batch_size = 32

        for start in tqdm(range(0, len(samples), batch_size), desc="Generating outputs"):
            batch_texts = samples[start : start + batch_size]
            batch_images = images[start : start + batch_size]

            model_inputs = self.processor(
                text=batch_texts,
                images=batch_images,
                return_tensors="pt",
                padding=True,
            ).to(self.device)

            # 3. Generate
            input_len = model_inputs["input_ids"].shape[-1]

            with torch.no_grad():
                output_ids = self.model.generate(
                    **model_inputs,
                    max_new_tokens=self.config.max_new_tokens,
                    do_sample=self.config.temperature > 0.0,
                    temperature=self.config.temperature,
                    top_p=self.config.top_p,
                )

            # 4. Slice and Decode
            # Slice off the input tokens to keep only the generated response
            generated_ids = output_ids[:, input_len:]
            decoded = self.processor.batch_decode(generated_ids, skip_special_tokens=True)
            outputs.extend([d.strip() for d in decoded])

        return outputs


class GemmaGenerator(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,
            limit_mm_per_prompt={"image": 1},
        )

        super().__init__(model)

    def _create_conversation(
        self, text: str, image: Union[Image, list[Image], None]
    ) -> list[dict]:
        conversation = []
        if self.config.system_prompt:
            conversation.append({"role": "system", "content": self.config.system_prompt})

        images = image if isinstance(image, list) else [image]
        user_content = [
            {"type": "image_pil", "image_pil": img() if callable(img) else img}
            for img in images
            if img is not None
        ]
        user_content.append({"type": "text", "text": text})

        conversation.append({"role": "user", "content": user_content})
        return conversation

    def _create_conversations(self, samples: list[tuple]) -> list[list[dict]]:
        conversations = []
        with ThreadPoolExecutor(max_workers=32) as executor:
            futures = []
            for text, image_obj in samples:
                futures.append(executor.submit(self._create_conversation, text, image_obj))

            for future in tqdm(futures, desc="Creating conversations"):
                conversations.append(future.result())

        return conversations

    def generate(
        self, samples: list[str], images: list[Union[Image, list[Image], None]]
    ) -> list[str]:
        sampling_params = SamplingParams(
            temperature=self.config.temperature,
            max_tokens=self.config.max_new_tokens,
            top_p=self.config.top_p,
        )

        outputs = []
        for conversations in batch_with_preload(
            items=list(zip(samples, images)),
            batch_size=self.config.batch_size,
            process_func=self._create_conversations,
        ):
            partial_outputs = self.model.chat(
                messages=conversations, sampling_params=sampling_params, use_tqdm=True
            )
            outputs.extend(partial_outputs)

        return [output.outputs[0].text for output in outputs]
