import base64
import os
from typing import Any

from inspect_ai.tool import Tool, ToolError, tool
from inspect_ai.util import sandbox
from openai import OpenAI


@tool
def query_vision_language_model(vllm_model: str = "gpt-4o-mini") -> Tool:
    client = OpenAI(
        api_key=os.environ.get("OPENAI_API_KEY"),
    )

    async def execute(prompt: str, input_image_paths: list[str]) -> Any:
        """Query a vision language model with an image and text input.

        Args:
            prompt (str): The prompt to use
            input_image_paths (list[str]): The path to the input image

        Returns:
            str: The response from the vision language model
        """
        base64_images = []
        for input_image_path in input_image_paths:
            # Read the image as base64
            try:
                result = await sandbox().read_file(input_image_path, text=False)
                base64_image = base64.b64encode(result).decode("utf-8")
            except (PermissionError, IsADirectoryError) as e:
                raise ToolError(str(e))
            base64_images.append(base64_image)

        # Construct the message with multiple images
        message_content: list[dict[str, Any]] = [{"type": "text", "text": prompt}]
        for base64_image in base64_images:
            message_content.append(
                {
                    "type": "image_url",
                    "image_url": {"url": f"data:image/jpeg;base64,{base64_image}"},
                }
            )

        # Payload for OpenAI API
        payload = {
            "model": vllm_model,
            "messages": [{"role": "user", "content": message_content}],
        }

        # Make a request to OpenAI API
        try:
            response = client.chat.completions.create(**payload)  # type: ignore
        except Exception as e:
            raise ToolError(str(e))
        return response.choices[0].message.content

    return execute
