"""Tests for the VisualWebBench utils module."""

from unittest.mock import patch

from screensuite.benchmarks.perception.visualwebbench.utils import (
    eval_element_ocr,
    eval_element_or_action,
    eval_heading_ocr_or_web_caption,
    eval_webqa,
)


def test_eval_web_caption():
    """Test the eval_web_caption function."""
    preds = ["This is a test caption", "Another test caption"]
    golds = ["This is a test caption", "Different caption"]

    with patch("screensuite.benchmarks.perception.visualwebbench.utils.Rouge") as mock_rouge:
        mock_rouge.return_value.get_scores.return_value = {
            "rouge-1": {"f": 0.8},
            "rouge-2": {"f": 0.7},
            "rouge-l": {"f": 0.75},
        }

        result = eval_heading_ocr_or_web_caption(preds, golds)

        # Check that Rouge was called with the correct arguments
        mock_rouge.assert_called_once_with(metrics=["rouge-1", "rouge-2", "rouge-l"])
        mock_rouge.return_value.get_scores.assert_called_once_with(preds, golds, avg=True)

        # Check the result
        expected_result = (0.8 + 0.7 + 0.75) / 3
        assert result == expected_result


def test_eval_element_ocr():
    """Test the eval_element_ocr function."""
    preds = ["This is a test element", "Another test element"]
    golds = ["This is a test element", "Different element"]

    with patch("screensuite.benchmarks.perception.visualwebbench.utils.Rouge") as mock_rouge:
        mock_rouge.return_value.get_scores.return_value = {
            "rouge-1": {"f": 0.8},
            "rouge-2": {"f": 0.7},
            "rouge-l": {"f": 0.75},
        }

        result = eval_element_ocr(preds, golds)

        # Check that Rouge was called with the correct arguments
        mock_rouge.assert_called_once_with(metrics=["rouge-1", "rouge-2", "rouge-l"])
        mock_rouge.return_value.get_scores.assert_called_once_with(preds, golds, avg=True)

        # Check the result
        expected_result = (0.8 + 0.7 + 0.75) / 3
        assert result == expected_result


def test_eval_element_or_action():
    """Test the eval_action_prediction function."""
    preds = ["A", "B", "C", "A"]
    golds = [0, 1, 3, 0]

    result = eval_element_or_action(preds, golds)

    # Check the result (3 out of 4 correct)
    assert result == 0.75


def test_eval_webqa():
    """Test the eval_webqa function."""
    preds = ["This is a test answer", "Another test answer"]
    golds = [["This is a test answer", "Different answer"]]

    with patch("screensuite.benchmarks.perception.visualwebbench.utils.Rouge") as mock_rouge:
        mock_rouge.return_value.get_scores.return_value = {
            "rouge-1": {"f": 0.8},
        }

        result = eval_webqa(preds, golds)

        # Check that Rouge was called with the correct arguments
        mock_rouge.assert_called_with(metrics=["rouge-1"])
        # Since we have 2 gold answers for the first prediction, get_scores should be called twice
        assert mock_rouge.return_value.get_scores.call_count == 2

        # Check the result
        assert result == 0.8


def test_empty_predictions():
    """Test evaluation functions with empty predictions."""
    preds = ["", "", ""]
    golds = ["Test 1", "Test 2", "Test 3"]

    # All evaluation functions should handle empty predictions
    with patch("screensuite.benchmarks.perception.visualwebbench.utils.Rouge") as mock_rouge:
        mock_rouge.return_value.get_scores.return_value = {
            "rouge-1": {"f": 0.0},
            "rouge-2": {"f": 0.0},
            "rouge-l": {"f": 0.0},
        }

        # These functions should not raise exceptions
        eval_heading_ocr_or_web_caption(preds, golds)
        eval_element_ocr(preds, golds)
        eval_webqa(preds, [golds])
