from unittest.mock import AsyncMock, patch

import pytest
from inspect_ai.model import ChatMessageAssistant, ChatMessageUser
from inspect_ai.scorer import Score
from inspect_ai.solver import TaskState
from inspect_ai.util import ExecResult

from inspect_evals.scicode.scorer import reduce_subproblems, verify, verify_subproblem


@pytest.fixture
def successful_subproblem():
    return {
        "step_number": "1.1",
        "test_cases": ["assert 1 + 1 == 2"],
    }


@pytest.fixture
def failing_subproblem():
    return {
        "step_number": "1.2",
        "test_cases": ["assert 1 + 1 == 3"],
    }


@pytest.fixture
def dummy_task_state(successful_subproblem, failing_subproblem):
    return TaskState(
        model="",
        sample_id="1",
        epoch=1,
        input="",
        messages=[
            ChatMessageUser(content="Solve 1.1"),
            ChatMessageAssistant(content="```python```"),
            ChatMessageUser(content="Solve 1.2"),
            ChatMessageAssistant(content="```python```"),
        ],
        metadata={
            "required_dependencies": "import numpy as np\n",
            "sub_steps": [successful_subproblem, failing_subproblem],
        },
    )


@pytest.fixture
def mock_sandbox():
    def _mock(result, times_out=False):
        mock_sandbox = AsyncMock()
        mock_sandbox.exec = AsyncMock(return_value=result)
        if times_out:
            mock_sandbox.exec.side_effect = TimeoutError("Execution timed out.")
        return mock_sandbox

    return _mock


@pytest.fixture
def sandbox_return_success():
    return ExecResult(
        success=True,
        returncode=0,
        stdout="subproblem passed",
        stderr="",
    )


@pytest.fixture
def sandbox_return_failure():
    return ExecResult(
        success=False,
        returncode=1,
        stdout="subproblem failed",
        stderr="ExampleError",
    )


@pytest.mark.asyncio
async def test_verify_subproblem_success(
    mock_sandbox, dummy_task_state, successful_subproblem, sandbox_return_success
):
    with patch(
        "inspect_evals.scicode.scorer.sandbox",
        return_value=mock_sandbox(sandbox_return_success),
    ):
        scorer = verify_subproblem(successful_subproblem, timeout=5)
        score = await scorer(dummy_task_state, None)
        # Mock execution was successful, so subproblem should be scored as 1.0
        assert score.value == {"1.1": 1.0}
        # Check that dependencies were included in "executed" code.
        assert "import numpy as np" in score.answer
        # Check that tests were included in "executed" code.
        assert "assert 1 + 1 == 2" in score.answer
        # Check that logic to unpack test case targets was included in "executed" code.
        assert "targets = process_hdf5_to_tuple" in score.answer
        # Explanation indicates success.
        assert "All test cases passed" in score.explanation


@pytest.mark.asyncio
async def test_verify_subproblem_failure(
    mock_sandbox, dummy_task_state, failing_subproblem, sandbox_return_failure
):
    with patch(
        "inspect_evals.scicode.scorer.sandbox",
        return_value=mock_sandbox(sandbox_return_failure),
    ):
        scorer = verify_subproblem(failing_subproblem, timeout=5)
        score = await scorer(dummy_task_state, None)
        # Mock execution failed, so subproblem should be scored as 0.0
        assert score.value == {"1.2": 0.0}
        # Check that dependencies were included in "executed" code.
        assert "import numpy as np" in score.answer
        # Check that tests were included in "executed" code.
        assert "assert 1 + 1 == 3" in score.answer
        # Check that logic to unpack test case targets was included in "executed" code.
        assert "targets = process_hdf5_to_tuple" in score.answer
        # Explanation indicates failure.
        assert "Code did not pass all test cases" in score.explanation
        # Error is included in explanation
        assert "ExampleError" in score.explanation


@pytest.mark.asyncio
async def test_verify_subproblem_timeout(
    mock_sandbox, dummy_task_state, successful_subproblem, sandbox_return_failure
):
    with patch(
        "inspect_evals.scicode.scorer.sandbox",
        return_value=mock_sandbox(sandbox_return_failure, times_out=True),
    ):
        scorer = verify_subproblem(successful_subproblem, timeout=5)
        score = await scorer(dummy_task_state, None)
        # Mock execution timed out, so subproblem should be scored as 0.0
        assert score.value == {"1.1": 0.0}
        # Check that dependencies were included in "executed" code.
        assert "import numpy as np" in score.answer
        # Check that tests were included in "executed" code.
        assert "assert 1 + 1 == 2" in score.answer
        # Check that logic to unpack test case targets was included in "executed" code.
        assert "targets = process_hdf5_to_tuple" in score.answer
        # Explanation indicates timeout.
        assert "execution timed out" in score.explanation


def test_reducer():
    scores = [
        Score(value={"1.1": 1.0}, explanation="1.1 passed."),
        Score(value={"1.2": 0.0}, explanation="1.2 failed."),
    ]

    reducer = reduce_subproblems()
    result = reducer(scores)

    assert result.value == {"1.1": 1.0, "1.2": 0.0}
    assert "1.1 passed." in result.explanation
    assert "1.2 failed." in result.explanation


@pytest.mark.asyncio
async def test_verify(
    mock_sandbox, dummy_task_state, sandbox_return_success, sandbox_return_failure
):
    with patch("inspect_evals.scicode.scorer.sandbox") as sandbox_patch:
        sandbox_mock = AsyncMock()
        sandbox_mock.exec = AsyncMock(
            side_effect=[sandbox_return_success, sandbox_return_failure]
        )
        sandbox_patch.return_value = sandbox_mock

        scorer = verify(timeout=5)
        score = await scorer(dummy_task_state, None)

        assert score.value == {"1.1": 1.0, "1.2": 0.0}

        explanation = score.explanation
        assert "All test cases passed for subproblem 1.1" in explanation
        assert "Code did not pass all test cases." in explanation
        assert "ExampleError" in explanation
