"""
QwenVisionProcessor
--------------------
Wraps CLIPImageProcessor (for images) + Qwen tokenizer (for text) into a
single AutoProcessor-compatible class.

Supports apply_chat_template() so callers can use the exact same interface
as granite-vision or LLaVA.
"""

import os
import requests
from io import BytesIO
from typing import List, Optional, Union
from PIL import Image

from transformers import (
    ProcessorMixin,
    CLIPImageProcessor,
    AutoTokenizer,
    BatchEncoding,
)


IMG_TOKEN = "[IMG]"
IMG_TOKEN_COUNT = 32


class QwenVisionProcessor(ProcessorMixin):
    """
    Processor for QwenVisionForConditionalGeneration.

    Attributes exposed for AutoProcessor
    -------------------------------------
    attributes = ["image_processor", "tokenizer"]
    """

    # Required by ProcessorMixin / AutoProcessor registry
    attributes = ["image_processor", "tokenizer"]
    image_processor_class = "CLIPImageProcessor"
    tokenizer_class = "AutoTokenizer"

    def __init__(self, image_processor: CLIPImageProcessor, tokenizer):
        super().__init__(image_processor, tokenizer)
        self.image_processor = image_processor
        self.tokenizer = tokenizer

        # Ensure [IMG] token exists
        if tokenizer.convert_tokens_to_ids(IMG_TOKEN) == tokenizer.unk_token_id:
            tokenizer.add_tokens([IMG_TOKEN])

        self.img_token = IMG_TOKEN
        self.img_token_id = tokenizer.convert_tokens_to_ids(IMG_TOKEN)
        self.img_token_count = IMG_TOKEN_COUNT

    # ------------------------------------------------------------------ #
    #  Factory methods                                                     #
    # ------------------------------------------------------------------ #

    @classmethod
    def from_pretrained(cls, pretrained_model_name_or_path: str, **kwargs):
        image_processor = CLIPImageProcessor.from_pretrained(
            "openai/clip-vit-base-patch32"
        )
        tokenizer = AutoTokenizer.from_pretrained(
            pretrained_model_name_or_path, **kwargs
        )
        return cls(image_processor=image_processor, tokenizer=tokenizer)

    def save_pretrained(self, save_directory: str, **kwargs):
        os.makedirs(save_directory, exist_ok=True)
        self.image_processor.save_pretrained(save_directory)
        self.tokenizer.save_pretrained(save_directory)

    # ------------------------------------------------------------------ #
    #  apply_chat_template  — mirrors the granite-vision interface        #
    # ------------------------------------------------------------------ #

    def apply_chat_template(
        self,
        conversation: List[dict],
        add_generation_prompt: bool = True,
        tokenize: bool = True,
        return_dict: bool = True,
        return_tensors: Optional[str] = "pt",
        images: Optional[List[Image.Image]] = None,
        max_length: int = 512,
        padding: Union[bool, str] = True,
        truncation: bool = True,
        enable_thinking: bool = False,
        **kwargs,
    ) -> BatchEncoding:
        """
        Parameters
        ----------
        conversation : list of dicts
            Each dict has "role" and "content".
            Content can be a string, or a list of dicts with "type" keys:
              {"type": "image", "url": "/path/to/img.png"}
              {"type": "text",  "text": "Your question"}
        images : optional pre-loaded PIL images (overrides url extraction)
        """

        # ---- Extract images from conversation ------------------------ #
        extracted_images: List[Image.Image] = []
        text_messages = []

        for turn in conversation:
            role = turn["role"]
            content = turn["content"]

            if isinstance(content, str):
                text_messages.append({"role": role, "content": content})
                continue

            # List of content blocks
            text_parts = []
            for block in content:
              if block.get("type") == "image":
                  if images is None:
                      url = block.get("url") or block.get("path")
                      if url:
                          if url.startswith("http://") or url.startswith("https://"):
                              response = requests.get(url)
                              extracted_images.append(Image.open(BytesIO(response.content)).convert("RGB"))
                          else:
                              extracted_images.append(Image.open(url).convert("RGB"))
                  img_placeholder = " ".join([self.img_token] * self.img_token_count)
                  text_parts.append(img_placeholder)
              elif block.get("type") == "text":
                  text_parts.append(block["text"])

            text_messages.append({"role": role, "content": " ".join(text_parts)})

        # If caller provided images explicitly, use those
        if images is not None:
            extracted_images = images

        # ---- Build prompt string via tokenizer's chat template ------- #
        prompt_text = self.tokenizer.apply_chat_template(
            text_messages,
            tokenize=False,
            add_generation_prompt=add_generation_prompt,
            enable_thinking=enable_thinking,
        )

        if not tokenize:
            return prompt_text  # type: ignore

        # ---- Tokenise text ------------------------------------------- #
        encoding = self.tokenizer(
            prompt_text,
            return_tensors=return_tensors,
            padding=padding,
            truncation=truncation,
            max_length=max_length,
            add_special_tokens=False,
        )

        # ---- Process images ------------------------------------------ #
        if extracted_images:
            pixel_values = self.image_processor(
                images=extracted_images, return_tensors=return_tensors
            )["pixel_values"]
            encoding["pixel_values"] = pixel_values
        else:
            # No image supplied — caller must add pixel_values separately
            pass

        if return_dict:
            return BatchEncoding(encoding)
        return encoding

    # ------------------------------------------------------------------ #
    #  Standard __call__                                                   #
    # ------------------------------------------------------------------ #

    def __call__(
        self,
        text: Optional[Union[str, List[str]]] = None,
        images: Optional[Union[Image.Image, List[Image.Image]]] = None,
        return_tensors: Optional[str] = "pt",
        padding: Union[bool, str] = True,
        truncation: bool = True,
        max_length: int = 512,
        **kwargs,
    ) -> BatchEncoding:

        encoding = {}

        if text is not None:
            text_enc = self.tokenizer(
                text,
                return_tensors=return_tensors,
                padding=padding,
                truncation=truncation,
                max_length=max_length,
                **kwargs,
            )
            encoding.update(text_enc)

        if images is not None:
            if isinstance(images, Image.Image):
                images = [images]
            pixel_values = self.image_processor(
                images=images, return_tensors=return_tensors
            )["pixel_values"]
            encoding["pixel_values"] = pixel_values

        return BatchEncoding(encoding)

    def decode(self, *args, **kwargs):
        return self.tokenizer.decode(*args, **kwargs)

    def batch_decode(self, *args, **kwargs):
        return self.tokenizer.batch_decode(*args, **kwargs)