"""Tests for WebSrc benchmark."""

from unittest.mock import patch

import pytest
from smolagents import ChatMessage

from screensuite.basebenchmark import EvaluationConfig
from screensuite.benchmark_result import BenchmarkResult
from screensuite.benchmarks.perception.websrc.benchmark import WebSrcBenchmark
from screensuite.benchmarks.perception.websrc.config import WebSrcConfig
from screensuite.benchmarks.perception.websrc.prompt import WebSrcPrompt


class MockModel:
    def __init__(self, responses):
        self.responses = responses

    def generate(self, messages, max_tokens, temperature):
        return ChatMessage(role="assistant", content=self.responses)


@pytest.fixture
def mock_dataset_1():
    return [
        {
            "question": "What is the color of the button?",
            "image": "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAIAAACQd1PeAAAADElEQVR4nGP4z8AAAAMBAQDJ/pLvAAAAAElFTkSuQmCC",
            "answer": "blue navy",
        },
    ]


@pytest.fixture
def mock_dataset_2():
    return [
        {
            "question": "What is the text on the screen?",
            "image": "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAIAAACQd1PeAAAADElEQVR4nGP4z8AAAAMBAQDJ/pLvAAAAAElFTkSuQmCC",
            "answer": "Hello World",
        },
    ]


@pytest.fixture
def websrc_benchmark():
    config = WebSrcConfig(
        hf_repo="rootsautomation/websrc",
        split="dev",
        system_prompt=WebSrcPrompt.SHORT_PROMPT,
        max_tokens=100,
        temperature=0.0,
    )
    return WebSrcBenchmark(name="test_websrc", config=config, tags=["test"])


def test_websrc_benchmark_initialization(websrc_benchmark):
    assert websrc_benchmark.name == "test_websrc"
    assert "test" in websrc_benchmark.tags
    assert websrc_benchmark.config.max_tokens == 100
    assert websrc_benchmark.config.temperature == 0.0


@patch("screensuite.benchmarks.perception.websrc.benchmark.HubBaseBenchmark.load")
def test_websrc_benchmark_evaluate(mock_load, websrc_benchmark, mock_dataset_1):
    # Setup
    websrc_benchmark.dataset = mock_dataset_1
    model = MockModel(responses="blue")

    # Execute
    results = websrc_benchmark.evaluate(
        model=model, evaluation_config=EvaluationConfig(timeout=10, parallel_workers=1, run_name=None)
    )

    # Assert
    assert isinstance(results, BenchmarkResult)
    assert "exact_match" in results.get_fields()
    assert "f1" in results.get_fields()
    assert isinstance(results["exact_match"], float)
    assert isinstance(results["f1"], float)
    assert results["exact_match"] == 0.0
    assert results["f1"] >= 0.66660


@patch("screensuite.benchmarks.perception.websrc.benchmark.HubBaseBenchmark.load")
def test_websrc_benchmark_evaluate_no_answer(mock_load, websrc_benchmark, mock_dataset_1):
    websrc_benchmark.dataset = mock_dataset_1
    model = MockModel(responses="Hello World")

    results = websrc_benchmark.evaluate(
        model=model, evaluation_config=EvaluationConfig(timeout=10, parallel_workers=1, run_name=None)
    )

    # First answer should be marked as incorrect since ground truth doesn't contain NO_ANSWER
    assert results["exact_match"] == 0.0
    assert results["f1"] == 0.0


@patch("screensuite.benchmarks.perception.websrc.benchmark.HubBaseBenchmark.load")
def test_websrc_benchmark_evaluate_empty_response(mock_load, websrc_benchmark, mock_dataset_2):
    # Setup
    websrc_benchmark.dataset = mock_dataset_2
    model = MockModel(responses=" Hello World")

    results = websrc_benchmark.evaluate(
        model=model, evaluation_config=EvaluationConfig(timeout=10, parallel_workers=1, run_name=None)
    )
    assert results["exact_match"] == 1.0
    assert results["f1"] == 1.0
