import json
import os
from unittest.mock import patch

from PIL import Image

from screensuite.basebenchmark import EvaluationConfig
from screensuite.benchmarks.singlestep.androidcontrol.benchmark import (
    AndroidControlBenchmark,
    AndroidControlConfig,
)
from screensuite.benchmarks.singlestep.androidcontrol.models import Operation


def create_empty_image():
    return Image.new("RGB", (100, 100), color="white")


@patch("PIL.Image.open")
def test_get_annotated_input_from_sample(mock_image_open):
    """Test the _get_annotated_input_from_sample method using the sample file."""
    # Setup mock to return empty image
    mock_image_open.return_value = create_empty_image()

    # Load the sample data
    current_dir = os.path.dirname(os.path.abspath(__file__))
    sample_file = os.path.join(current_dir, "sample.json")
    sample = json.load(open(sample_file))
    # Create benchmark instance
    config = AndroidControlConfig(hf_repo="smolagents/android-control", max_step_memory=3)
    benchmark = AndroidControlBenchmark("test_benchmark", config, ["test"])

    # Transform the sample
    transformed_samples = list(benchmark._get_annotated_input_from_sample(sample, EvaluationConfig(run_name="test")))

    # Verify the transformed samples
    assert len(transformed_samples) == 2

    for annotated_input in transformed_samples:
        query, (operations, target_dimensions) = annotated_input.messages, annotated_input.ground_truth
        # Verify query structure
        assert isinstance(query, list), "Query should be a list of messages"
        assert len(query) == 1
        message = query[0]
        assert isinstance(message, dict), "Message should be a dictionary"
        assert message["role"] == "user", "Message role should be user"

        # Verify content structure
        content = message["content"]
        assert len([c for c in content if c["type"] == "image"]) <= benchmark.config.max_step_memory + 1, (
            f"Message should have {benchmark.config.max_step_memory + 1} or less images"
        )

        # Verify operations and bounding box
        assert isinstance(operations, list), "Operations should be a list"
        assert all(isinstance(operation, Operation) for operation in operations), (
            "All operations should be Operation instances"
        )
        assert isinstance(target_dimensions, tuple)
        assert len(operations) == 1
