import base64
import io

import requests
import torch
from PIL import Image
from transformers import AutoModel, AutoProcessor

DEFAULT_PROMPT = "Please output the text content from the image."


class EndpointHandler:
    def __init__(self, path=""):
        self.processor = AutoProcessor.from_pretrained(path, trust_remote_code=True)
        if not torch.cuda.is_available():
            raise RuntimeError(
                "No GPU found. Recreate the endpoint on a GPU instance (e.g. Nvidia L4 or T4)."
            )
        self.model = (
            AutoModel.from_pretrained(
                path,
                trust_remote_code=True,
                torch_dtype=torch.bfloat16,
            )
            .to("cuda")
            .eval()
        )

    def _load_image(self, image):
        # Accept a URL, a base64 string, or a data URL ("data:image/png;base64,...")
        if image.startswith("http://") or image.startswith("https://"):
            raw = requests.get(image, timeout=30).content
        else:
            if "," in image and image.startswith("data:"):
                image = image.split(",", 1)[1]
            raw = base64.b64decode(image)
        return Image.open(io.BytesIO(raw)).convert("RGB")

    def __call__(self, data):
        inputs = data.get("inputs", data)
        image = self._load_image(inputs["image"])
        prompt = inputs.get("prompt", DEFAULT_PROMPT)
        max_new_tokens = int(inputs.get("max_new_tokens", 4096))

        messages = [
            {"role": "system", "content": "You are a helpful assistant."},
            {"role": "user", "content": [
                {"type": "image"},
                {"type": "text", "text": prompt},
            ]},
        ]
        chat_prompt = self.processor.apply_chat_template(
            messages, tokenize=False, add_generation_prompt=True
        )
        model_inputs = self.processor(
            text=[chat_prompt],
            images=[image],
            padding=True,
            return_tensors="pt",
        ).to(device=self.model.device, dtype=self.model.dtype)

        with torch.no_grad():
            output_ids = self.model.generate(
                **model_inputs,
                use_cache=True,
                max_new_tokens=max_new_tokens,
                do_sample=False,
            )

        new_tokens = output_ids[0][model_inputs.input_ids.shape[1]:]
        text = self.processor.decode(
            new_tokens,
            skip_special_tokens=True,
            clean_up_tokenization_spaces=False,
        ).strip()
        return {"text": text}