"""Tests for the NIAH scoring module."""

import pytest
from inspect_ai.scorer import SampleScore, Score, Target
from inspect_ai.solver import TaskState

from inspect_evals.niah.utils.scoring import (
    custom_model_graded_qa_with_history_scorer,
    return_metadata_variable_as_history,
    subset_accuracy_combinations,
)


def test_subset_accuracy_combinations():
    """Test the subset accuracy combinations metric."""
    # Create some test scores with metadata
    scores = [
        SampleScore(
            score=Score(
                value=1.0,
                metadata={
                    "target_context_length": 1000,
                    "target_position": 500,
                },
            ),
        ),
        SampleScore(
            score=Score(
                value=0.0,
                metadata={
                    "target_context_length": 1000,
                    "target_position": 500,
                },
            ),
        ),
        SampleScore(
            score=Score(
                value=1.0,
                metadata={
                    "target_context_length": 2000,
                    "target_position": 1000,
                },
            ),
        ),
    ]

    # Calculate metrics
    metric_fn = subset_accuracy_combinations()
    results = metric_fn(scores)

    # Check individual parameter accuracies
    assert results["target_context_length_1000_accuracy"] == 0.5  # 1/2 correct
    assert results["target_context_length_2000_accuracy"] == 1.0  # 1/1 correct
    assert results["target_position_500_accuracy"] == 0.5  # 1/2 correct
    assert results["target_position_1000_accuracy"] == 1.0  # 1/1 correct

    # Check combined parameter accuracies
    assert (
        results["target_context_length_1000_target_position_500_accuracy"] == 0.5
    )  # 1/2 correct
    assert (
        results["target_context_length_2000_target_position_1000_accuracy"] == 1.0
    )  # 1/1 correct

    # Check overall accuracy
    assert results["overall_accuracy"] == 2 / 3  # 2/3 correct total


class MockTaskState(TaskState):
    def __init__(self, metadata):
        super().__init__(
            model="test_model",
            sample_id="test_sample",
            epoch=0,
            messages=[],
            input="test input",
            output=None,  # ModelOutput | None
            metadata=metadata,
        )


async def test_custom_scorer_wrapper() -> None:
    """Test the custom scorer wrapper adds metadata from the task state"""
    # Create a mock task state with metadata
    task_state = MockTaskState(
        metadata={"needle_question": "What is the test question?"},
    )

    # Create the wrapped scorer
    score_fn = custom_model_graded_qa_with_history_scorer(
        include_history=return_metadata_variable_as_history,
        judge_model="mockllm/model",
    )

    # Create a target
    target = Target("test target")

    # Get the score
    score: Score = await score_fn(task_state, target)

    # Verify that the score has metadata
    assert score.metadata == task_state.metadata
    assert isinstance(score.as_str(), str)
    assert isinstance(score.explanation, str)


def test_return_metadata_variable_as_history_valid_string():
    """Test return_metadata_variable_as_history with valid string metadata."""
    task_state = MockTaskState(
        metadata={"needle_question": "What is the test question?"},
    )

    result = return_metadata_variable_as_history(task_state)
    assert result == "What is the test question?"


def test_return_metadata_variable_as_history_missing_key():
    """Test return_metadata_variable_as_history with missing metadata key."""
    task_state = MockTaskState(
        metadata={"other_key": "value"},
    )

    with pytest.raises(
        KeyError, match="No variable with name `needle_question` found in metadata."
    ):
        return_metadata_variable_as_history(task_state)


def test_return_metadata_variable_as_history_non_string():
    """Test return_metadata_variable_as_history with non-string metadata value."""
    task_state = MockTaskState(
        metadata={"needle_question": 42},
    )

    result = return_metadata_variable_as_history(task_state)
    assert result == "42"
