import base64
import os
from unittest.mock import AsyncMock, MagicMock, patch

import pytest
from inspect_ai.tool import ToolError

from inspect_evals.core_bench.tools import query_vision_language_model


@patch("openai.OpenAI")
@patch("inspect_evals.core_bench.tools.sandbox")
@patch.dict(os.environ, {"OPENAI_API_KEY": "test_api_key"})
async def test_query_vision_language_model_success(
    mock_sandbox, mock_openai, mock_openai_client
):
    """Test successful query to vision language model"""
    # Setup mocks
    sandbox_instance = MagicMock()
    sandbox_instance.read_file = AsyncMock()
    sandbox_instance.read_file.return_value = (
        b"\xff\xd8\xff\xe0\x00\x10JFIF\x00\x01\x01"
    )
    mock_sandbox.return_value = sandbox_instance
    mock_openai.return_value = mock_openai_client

    # Get the tool function
    tool_fn = query_vision_language_model("gpt-4o-mini")

    # Execute the tool with test inputs
    result = await tool_fn(
        prompt="Describe this image",
        input_image_paths=["/path/to/image1.jpg", "/path/to/image2.jpg"],
    )

    # Check the result
    assert result == "This is a test response from the vision model"

    # Verify sandbox was called to read the files
    assert sandbox_instance.read_file.call_count == 2

    # Verify OpenAI client was configured correctly
    mock_openai.assert_called_once_with(api_key="test_api_key")

    # Verify OpenAI API was called with the correct payload
    mock_openai_client.chat.completions.create.assert_called_once()
    call_args = mock_openai_client.chat.completions.create.call_args[1]

    # Check model
    assert call_args["model"] == "gpt-4o-mini"

    # Check message structure
    messages = call_args["messages"]
    assert len(messages) == 1
    assert messages[0]["role"] == "user"

    # Check content structure (text + 2 images)
    content = messages[0]["content"]
    assert len(content) == 3
    assert content[0]["type"] == "text"
    assert content[0]["text"] == "Describe this image"
    assert content[1]["type"] == "image_url"
    assert content[2]["type"] == "image_url"

    # Check that base64 encoding was performed
    for i in range(1, 3):
        assert "data:image/jpeg;base64," in content[i]["image_url"]["url"]


@patch("inspect_evals.core_bench.tools.sandbox")
@patch.dict(os.environ, {"OPENAI_API_KEY": "test_api_key"})
async def test_query_vision_language_model_file_error(mock_sandbox):
    """Test handling of file read errors"""
    # Setup mock to raise an error
    sandbox_instance = MagicMock()
    sandbox_instance.read_file = AsyncMock()
    sandbox_instance.read_file.side_effect = PermissionError("Permission denied")
    mock_sandbox.return_value = sandbox_instance

    # Get the tool function
    tool_fn = query_vision_language_model()

    # Execute the tool and check for error
    with pytest.raises(ToolError):
        await tool_fn(
            prompt="Describe this image", input_image_paths=["/path/to/image.jpg"]
        )


@patch("openai.OpenAI")
@patch("inspect_evals.core_bench.tools.sandbox")
@patch.dict(os.environ, {"OPENAI_API_KEY": "test_api_key"})
async def test_query_vision_language_model_custom_model(
    mock_sandbox, mock_openai, mock_openai_client
):
    """Test using a custom model"""
    # Setup mocks
    sandbox_instance = MagicMock()
    sandbox_instance.read_file = AsyncMock()
    sandbox_instance.read_file.return_value = (
        b"\xff\xd8\xff\xe0\x00\x10JFIF\x00\x01\x01"
    )
    mock_sandbox.return_value = sandbox_instance
    mock_openai.return_value = mock_openai_client

    # Get the tool function with a custom model
    tool_fn = query_vision_language_model("gpt-4o")

    # Execute the tool
    await tool_fn(
        prompt="Describe this image", input_image_paths=["/path/to/image.jpg"]
    )

    # Verify the correct model was used
    call_args = mock_openai_client.chat.completions.create.call_args[1]
    assert call_args["model"] == "gpt-4o"


@patch("inspect_evals.core_bench.tools.sandbox")
@patch.dict(os.environ, {"OPENAI_API_KEY": "test_api_key"})
async def test_query_vision_language_model_directory_error(mock_sandbox):
    """Test handling of directory errors"""
    # Setup mock to raise a directory error
    sandbox_instance = MagicMock()
    sandbox_instance.read_file = AsyncMock()
    sandbox_instance.read_file.side_effect = IsADirectoryError("Is a directory")
    mock_sandbox.return_value = sandbox_instance

    # Get the tool function
    tool_fn = query_vision_language_model()

    # Execute the tool and check for error
    with pytest.raises(ToolError):
        await tool_fn(
            prompt="Describe this image", input_image_paths=["/path/to/directory"]
        )


@patch("openai.OpenAI")
@patch("inspect_evals.core_bench.tools.sandbox")
@patch.dict(os.environ, {"OPENAI_API_KEY": "test_api_key"})
async def test_base64_encoding(mock_sandbox, mock_openai, mock_openai_client):
    """Test that the image is properly base64 encoded"""
    # Setup mocks
    sandbox_instance = MagicMock()
    sandbox_instance.read_file = AsyncMock()
    # Setup mock to return a known binary value
    test_binary = b"test_image_data"
    sandbox_instance.read_file.return_value = test_binary
    mock_sandbox.return_value = sandbox_instance
    mock_openai.return_value = mock_openai_client

    # Get the tool function
    tool_fn = query_vision_language_model()

    # Execute the tool
    await tool_fn(
        prompt="Describe this image", input_image_paths=["/path/to/image.jpg"]
    )

    # Get the content from the OpenAI call
    call_args = mock_openai_client.chat.completions.create.call_args[1]
    content = call_args["messages"][0]["content"]
    image_url = content[1]["image_url"]["url"]

    # Extract the base64 part
    base64_part = image_url.split(",")[1]

    # Decode and verify
    decoded = base64.b64decode(base64_part)
    assert decoded == test_binary
