from typing import overload, Literal
import re
import base64
from io import BytesIO

from PIL import Image
import torch
from openai import OpenAI
from tenacity import retry, stop_after_attempt, wait_fixed


class Captioner:

    def __init__(self, device: torch.device) -> "Captioner":
        self.device = device

    @overload
    def __call__(self, image: Image.Image) -> str: ...


class EmptyCaptioner(Captioner):

    def __call__(self, image: Image.Image) -> str:
        return ""


class GPTCaptioner(Captioner):

    DEFAULT_PROMPT = "Provide a detailed description of this image without exceeding 100 words and without line breaks."

    def __init__(self, api_key, base_url, model):
        self.client = OpenAI(api_key=api_key, base_url=base_url)
        self.model = model

    @staticmethod
    def pil_image_to_base64(image, format="PNG"):
        buffered = BytesIO()
        image.save(buffered, format=format)
        img_bytes = buffered.getvalue()
        img_base64 = base64.b64encode(img_bytes).decode("utf-8")
        return img_base64

    @retry(stop=stop_after_attempt(5), wait=wait_fixed(1))
    def get_response(self, base64_image, prompt):
        response = self.client.chat.completions.create(
            model=self.model,
            messages=[
                {
                    "role": "user",
                    "content": [
                        {
                            "type": "text",
                            "text": prompt,
                        },
                        {
                            "type": "image_url",
                            "image_url": {
                                "url": f"data:image/png;base64,{base64_image}",
                                "detail": "high",
                            },
                        },
                    ],
                }
            ],
            stream=False,
        )
        return response.choices[0].message.content

    def __call__(self, image, prompt=None):
        base64_image = self.pil_image_to_base64(image)
        if prompt is None:
            prompt = self.DEFAULT_PROMPT
        caption = self.get_response(base64_image, prompt=prompt)
        return caption