"""Tests for the VisualWebBench benchmark."""

from unittest.mock import patch

import pytest
from PIL import Image
from smolagents.models import ChatMessage

from screensuite.basebenchmark import EvaluationConfig
from screensuite.benchmarks.perception.visualwebbench.benchmark import (
    VisualWebBenchBenchmark,
)
from screensuite.benchmarks.perception.visualwebbench.config import (
    VisualWebBenchConfig,
    VisualWebBenchTask,
)


class MockModel:
    def __init__(self):
        self.call_count = 0

    def generate(self, message, max_tokens, temperature):
        self.call_count += 1
        return ChatMessage(role="assistant", content="Here is the answer")


@pytest.fixture
def mock_dataset():
    """Create a mock dataset for testing."""
    return {
        VisualWebBenchTask.CAPTION_TASK.value: [
            {"image": Image.new("RGB", (1, 1), color=(255, 255, 255)), "answer": "Test caption 1"},
            {"image": Image.new("RGB", (1, 1), color=(255, 255, 255)), "answer": "Test caption 2"},
        ],
        VisualWebBenchTask.HEADING_OCR_TASK.value: [
            {"image": Image.new("RGB", (1, 1), color=(255, 255, 255)), "answer": "Test heading 1"},
            {"image": Image.new("RGB", (1, 1), color=(255, 255, 255)), "answer": "Test heading 2"},
        ],
        VisualWebBenchTask.WEBQA_TASK.value: [
            {
                "image": Image.new("RGB", (1, 1), color=(255, 255, 255)),
                "question": "Test question 1",
                "answer": "Test answer 1",
            },
            {
                "image": Image.new("RGB", (1, 1), color=(255, 255, 255)),
                "question": "Test question 2",
                "answer": "Test answer 2",
            },
        ],
        VisualWebBenchTask.ELEMENT_OCR_TASK.value: [
            {
                "image": Image.new("RGB", (1, 1), color=(255, 255, 255)),
                "bbox": "0.1, 0.2, 0.3, 0.4",
                "answer": "Test text 1",
            },
            {
                "image": Image.new("RGB", (1, 1), color=(255, 255, 255)),
                "bbox": "0.2, 0.3, 0.4, 0.5",
                "answer": "Test text 2",
            },
        ],
        VisualWebBenchTask.ELEMENT_GROUND_TASK.value: [
            {"image": Image.new("RGB", (1, 1), color=(255, 255, 255)), "elem_desc": "Test element 1", "answer": "A"},
            {"image": Image.new("RGB", (1, 1), color=(255, 255, 255)), "elem_desc": "Test element 2", "answer": "B"},
        ],
        VisualWebBenchTask.ACTION_PREDICTION_TASK.value: [
            {
                "image": Image.new("RGB", (1, 1), color=(255, 255, 255)),
                "bbox": "0.1, 0.2, 0.3, 0.4",
                "options": "A. Option 1\nB. Option 2",
                "answer": "A",
            },
            {
                "image": Image.new("RGB", (1, 1), color=(255, 255, 255)),
                "bbox": "0.2, 0.3, 0.4, 0.5",
                "options": "A. Option 1\nB. Option 2",
                "answer": "B",
            },
        ],
        VisualWebBenchTask.ACTION_GROUND_TASK.value: [
            {
                "image": Image.new("RGB", (1, 1), color=(255, 255, 255)),
                "instruction": "Test instruction 1",
                "answer": "A",
            },
            {
                "image": Image.new("RGB", (1, 1), color=(255, 255, 255)),
                "instruction": "Test instruction 2",
                "answer": "B",
            },
        ],
    }


@pytest.fixture
def mock_metrics():
    """Create mock metrics for testing."""
    return {
        VisualWebBenchTask.CAPTION_TASK: lambda preds, golds: 0.8,
        VisualWebBenchTask.HEADING_OCR_TASK: lambda preds, golds: 0.7,
        VisualWebBenchTask.WEBQA_TASK: lambda preds, golds: 0.6,
        VisualWebBenchTask.ELEMENT_OCR_TASK: lambda preds, golds: 0.5,
        VisualWebBenchTask.ELEMENT_GROUND_TASK: lambda preds, golds: 0.4,
        VisualWebBenchTask.ACTION_PREDICTION_TASK: lambda preds, golds: 0.3,
        VisualWebBenchTask.ACTION_GROUND_TASK: lambda preds, golds: 0.2,
    }


@pytest.fixture
def benchmark():
    """Create a VisualWebBenchBenchmark instance for testing."""
    return VisualWebBenchBenchmark(
        name="test-visualwebbench",
        config=VisualWebBenchConfig.create(),
        tags=["test", "visualwebbench"],
    )


def test_prompt_formatters(benchmark):
    """Test that all task types have prompt formatters."""
    for task_type in VisualWebBenchTask:
        assert task_type in benchmark._prompt_formatters
        assert callable(benchmark._prompt_formatters[task_type])


def test_metrics(benchmark):
    """Test that all task types have metrics."""
    for task_type in VisualWebBenchTask:
        assert task_type in benchmark._metrics
        assert callable(benchmark._metrics[task_type])


def test_process_response_caption_task(benchmark):
    """Test processing response for caption task."""
    response = '<meta name="description" content="Test description">'
    processed = benchmark._process_response(response, VisualWebBenchTask.CAPTION_TASK)
    assert processed == "Test description"

    # Test with no meta tag
    response = "Some other content"
    processed = benchmark._process_response(response, VisualWebBenchTask.CAPTION_TASK)
    assert processed == "Some other content"

    # Test with None response
    processed = benchmark._process_response(None, VisualWebBenchTask.CAPTION_TASK)
    assert processed == ""


def test_process_response_action_prediction_task(benchmark):
    """Test processing response for action prediction task."""
    response = "a"
    processed = benchmark._process_response(response, VisualWebBenchTask.ACTION_PREDICTION_TASK)
    assert processed == "A"

    # Test with empty response
    response = ""
    processed = benchmark._process_response(response, VisualWebBenchTask.ACTION_PREDICTION_TASK)
    assert processed == ""

    # Test with None response
    processed = benchmark._process_response(None, VisualWebBenchTask.ACTION_PREDICTION_TASK)
    assert processed == ""


def test_process_response_webqa_task(benchmark):
    """Test processing response for webqa task."""
    response = "Answer: Test answer"
    processed = benchmark._process_response(response, VisualWebBenchTask.WEBQA_TASK)
    assert processed == "Test answer"

    # Test with no colon
    response = "Test answer"
    processed = benchmark._process_response(response, VisualWebBenchTask.WEBQA_TASK)
    assert processed == "Test answer"

    # Test with None response
    processed = benchmark._process_response(None, VisualWebBenchTask.WEBQA_TASK)
    assert processed == ""


def test_process_response_element_ocr_task(benchmark):
    """Test processing response for element OCR task."""
    response = "The text content within the red bounding box is: Test text"
    processed = benchmark._process_response(response, VisualWebBenchTask.ELEMENT_OCR_TASK)
    assert processed == "Test text"

    # Test with no colon
    response = "Test text"
    processed = benchmark._process_response(response, VisualWebBenchTask.ELEMENT_OCR_TASK)
    assert processed == "Test text"

    # Test with None response
    processed = benchmark._process_response(None, VisualWebBenchTask.ELEMENT_OCR_TASK)
    assert processed == ""


@patch("screensuite.benchmarks.perception.visualwebbench.benchmark.HubBaseBenchmark.load")
def test_evaluate(mock_load, benchmark, mock_dataset, mock_metrics):
    """Test the evaluate method."""
    # Set up the benchmark with mock data
    benchmark.datasets = mock_dataset
    benchmark._metrics = mock_metrics

    mock_model = MockModel()
    # Run the evaluation
    results = benchmark.evaluate(mock_model, EvaluationConfig(test_mode=False, parallel_workers=2, run_name=None))

    # Check the results
    assert results["web_caption"] == 0.8
    assert results["heading_ocr"] == 0.7
    assert results["webqa"] == 0.6
    assert results["element_ocr"] == 0.5
    assert results["element_ground"] == 0.4
    assert results["action_prediction"] == 0.3
    assert results["action_ground"] == 0.2

    # Check the average
    expected_average = (0.8 + 0.7 + 0.6 + 0.5 + 0.4 + 0.3 + 0.2) / 7
    assert results["average"] == expected_average

    # Check that the model was called for each sample
    expected_calls = sum(len(samples) for samples in mock_dataset.values())
    assert mock_model.call_count == expected_calls


@patch("screensuite.benchmarks.perception.visualwebbench.benchmark.HubBaseBenchmark.load")
def test_evaluate_with_missing_dataset(mock_load, benchmark, mock_dataset, mock_metrics):
    """Test the evaluate method when dataset is None."""
    benchmark.dataset = None
    mock_load.return_value = None

    # Set up the benchmark with mock data
    benchmark.datasets = {VisualWebBenchTask.CAPTION_TASK.value: mock_dataset[VisualWebBenchTask.CAPTION_TASK.value]}
    benchmark._metrics = {VisualWebBenchTask.CAPTION_TASK: mock_metrics[VisualWebBenchTask.CAPTION_TASK]}

    # Run the evaluation
    with pytest.raises(KeyError):
        results = benchmark.evaluate(MockModel(), EvaluationConfig(test_mode=False, parallel_workers=1, run_name=None))
