from typing import Dict, Any
from transformers import QwenImageProcessor, QwenTokenizer, QwenForMultiModalConditionalGeneration
import torch
from PIL import Image
import io
import json
import base64
import requests

class EndpointHandler():
    def __init__(self, path=""):
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        self.model = QwenForMultiModalConditionalGeneration.from_pretrained(
            path,
            torch_dtype=torch.float16 if self.device.type == "cuda" else torch.float32
        ).to(self.device)
        self.image_processor = QwenImageProcessor.from_pretrained(path)
        self.tokenizer = QwenTokenizer.from_pretrained(path)
        self.model.generation_config.use_cache = False 

    def __call__(self, data: Any) -> Dict[str, Any]:
        """
        Args:
            data (Any): The input data, which can be:
                - Binary image data in the request body.
                - A dictionary with 'image' and 'text' keys:
                    - 'image': Base64-encoded image string or image URL.
                    - 'text': The text prompt.

        Returns:
            Dict[str, Any]: The generated text output from the model.
        """
        if isinstance(data, (bytes, bytearray)):
            image = Image.open(io.BytesIO(data)).convert('RGB')
            text_input = "<|im_start|>user\nDescribe this image.\n<|im_end|><|im_start|>assistant\n"
        elif isinstance(data, dict):
            image_input = data.get('image', None)
            text_input = data.get('text', '')
            if image_input is None:
                return {"error": "No image provided."}
            if image_input.startswith('http'):
                response = requests.get(image_input)
                image = Image.open(io.BytesIO(response.content)).convert('RGB')
            else:
                image_data = base64.b64decode(image_input)
                image = Image.open(io.BytesIO(image_data)).convert('RGB')
        else:
            return {"error": "Invalid input data. Expected binary image data or a dictionary with 'image' key."}

        image_inputs = self.image_processor(images=image, return_tensors="pt").to(self.device)

        if not text_input:
            text_input = "<|im_start|>user\nDescribe this image.\n<|im_end|><|im_start|>assistant\n"
        input_ids = self.tokenizer(text_input, return_tensors="pt").input_ids.to(self.device)

        generated_ids = self.model.generate(
            **image_inputs,
            input_ids=input_ids,
            max_new_tokens=256,
            do_sample=True,
            top_p=0.9,
            temperature=0.7,
        )
        output_text = self.tokenizer.decode(generated_ids[0], skip_special_tokens=True)

        return {"generated_text": output_text}