"""Tests for ScreenQA utility functions."""

from screensuite.benchmarks.perception.screenqa.utils import (
    NO_ANSWER,
    f1_score,
    normalize_squad,
    sqa_s_metrics,
)


def test_normalize_squad():
    # Test basic normalization
    assert normalize_squad("The Quick Brown Fox") == "quick brown fox"

    # Test with punctuation
    assert normalize_squad("Hello, World!") == "hello world"

    # Test with extra whitespace
    assert normalize_squad("multiple    spaces") == "multiple spaces"

    # Test with articles
    assert normalize_squad("a cat and the dog") == "cat and dog"


def test_f1_score():
    # Test exact match
    pred = ["hello", "world"]
    gold = ["hello", "world"]
    assert f1_score(pred, gold) == 1.0

    # Test partial match
    pred = ["hello", "world"]
    gold = ["hello", "there"]
    expected_f1 = 0.5
    assert f1_score(pred, gold) == expected_f1

    # Test no match
    pred = ["hello", "world"]
    gold = ["goodbye", "world"]
    expected_f1 = 0.5
    assert f1_score(pred, gold) == expected_f1

    # Test empty lists
    assert f1_score([], []) == 0
    assert f1_score(["hello"], []) == 0
    assert f1_score([], ["hello"]) == 0


def test_sqa_s_metrics():
    # Test exact match
    pred = "Hello World"
    gold = ["Hello World"]
    exact_match, f1 = sqa_s_metrics(pred, gold)
    assert exact_match == 1
    assert f1 == 1.0

    # Test NO_ANSWER cases
    exact_match, f1 = sqa_s_metrics(NO_ANSWER, [NO_ANSWER])
    assert exact_match == 1
    assert f1 == 1.0

    exact_match, f1 = sqa_s_metrics(NO_ANSWER, ["Some Answer"])
    assert exact_match == 0
    assert f1 == 0.0

    # Test multiple ground truths
    pred = "Hello World"
    gold = ["Hello World", "Hello There"]
    exact_match, f1 = sqa_s_metrics(pred, gold)
    assert exact_match == 1
    assert f1 == 1.0

    # Test partial match
    pred = "Hello World"
    gold = ["Hello There"]
    exact_match, f1 = sqa_s_metrics(pred, gold)
    assert exact_match == 0
    assert f1 < 1.0
