import base64
import time
from mimetypes import guess_type
from openai import AzureOpenAI


class GPT4O():
    def __init__(self, api_key, api_version, deployment_name, base_url):
        self.api_key = api_key
        self.api_version = api_version
        self.deployment_name = deployment_name
        self.base_url = base_url
        self.client = AzureOpenAI(
            api_key=self.api_key,
            api_version=self.api_version,
            base_url=self.base_url
        )

    def local_image_to_data_url(self, image_path):
        mime_type, _ = guess_type(image_path)
        if mime_type is None:
            mime_type = 'application/octet-stream'
        with open(image_path, "rb") as image_file:
            base64_encoded = base64.b64encode(image_file.read()).decode('utf-8')
        return f"data:{mime_type};base64,{base64_encoded}"

    def chat(self, user_prompt, image_paths, max_retries=5):
        messages = [{
            "role": "user",
            "content": [{"type": "text", "text": user_prompt}]
        }]
        for image_path in image_paths:
            messages[0]["content"].append({
                "type": "image_url",
                "image_url": {
                    "url": self.local_image_to_data_url(image_path)
                }
            })
        for attempt in range(max_retries):
            try:
                response = self.client.chat.completions.create(
                    model=self.deployment_name,
                    messages=messages
                )
                return response.choices[0].message.content
            except Exception as e:
                print(f"[GPT-4o] Error attempt {attempt+1}: {e}")
                time.sleep(1)  
        return None 