import os
from typing import Any

from qwen_vl_utils import process_vision_info
from smolagents.models import ChatMessage, MessageRole, Model
from transformers import AutoProcessor, Qwen2_5_VLForConditionalGeneration


class QwenVLModel(Model):
    """Model wrapper for Qwen2.5VL"""

    def __init__(self, model_path: str = "Qwen/Qwen2.5-VL-3B-Instruct", device: str = "auto"):
        super().__init__()
        self.model_path = model_path
        self.device = device
        self.model_id = model_path
        self.model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
            model_path, torch_dtype="auto", device_map=device
        )
        self.processor = AutoProcessor.from_pretrained(model_path)

    def __call__(
        self, messages: list[dict[str, Any]], stop_sequences: list[str] | None = None, **kwargs
    ) -> ChatMessage:
        """Convert a list of messages to a model input and run inference"""

        # Count images in messages - debug
        image_count = 0
        for msg in messages:
            if isinstance(msg.get("content"), list):
                for item in msg["content"]:
                    if isinstance(item, dict) and item.get("type") == "image":
                        image_count += 1

        print(f"QwenVLModel received {len(messages)} messages with {image_count} images")

        # Format the messages for Qwen2.5VL
        formatted_messages = []

        for msg in messages:
            role = msg["role"]
            if isinstance(msg["content"], list):
                content = []
                for item in msg["content"]:
                    if item["type"] == "text":
                        content.append({"type": "text", "text": item["text"]})
                    elif item["type"] == "image":
                        # Handle image path or direct image object
                        if isinstance(item["image"], str):
                            print("Image in memory")
                            content.append({"type": "image", "image": item["image"]})
                        else:
                            print("Image in temp file")
                            # Save temporary image file
                            temp_img_path = "temp_image.png"
                            item["image"].save(temp_img_path)
                            content.append({"type": "image", "image": temp_img_path})
            else:
                content = [{"type": "text", "text": msg["content"]}]

            formatted_messages.append({"role": role, "content": content})

        # Prepare for inference
        text = self.processor.apply_chat_template(formatted_messages, tokenize=False, add_generation_prompt=True)
        image_inputs, video_inputs = process_vision_info(formatted_messages)
        inputs = self.processor(
            text=[text],
            images=image_inputs,
            videos=video_inputs,
            padding=True,
            return_tensors="pt",
        )
        inputs = inputs.to(self.model.device)

        # Generate the response
        generated_ids = self.model.generate(**inputs, max_new_tokens=512)
        generated_ids_trimmed = [out_ids[len(in_ids) :] for in_ids, out_ids in zip(inputs.input_ids, generated_ids)]
        output_text = self.processor.batch_decode(
            generated_ids_trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False
        )[0]

        # Clean up any temporary files
        if os.path.exists("temp_image.png"):
            os.remove("temp_image.png")

        return ChatMessage(role=MessageRole.ASSISTANT, content=output_text)

    def to_dict(self) -> dict[str, Any]:
        """Convert the model to a dictionary"""
        return {"class": self.__class__.__name__, "model_path": self.model_path, "device": self.device}

    @classmethod
    def from_dict(cls, data: dict[str, Any]) -> "QwenVLModel":
        """Create a model from a dictionary"""
        return cls(model_path=data.get("model_path", "Qwen/Qwen2.5-VL-3B-Instruct"), device=data.get("device", "auto"))
