from unittest.mock import AsyncMock, patch

import pytest
from inspect_ai.scorer import CORRECT, INCORRECT, Target
from inspect_ai.solver import TaskState
from inspect_ai.util import ExecResult

from inspect_evals.apps.apps import find_code, record_to_sample, verify


def test_find_code_with_markdown_block_and_text_after():
    completion = """
```python
def solution(x):
    return x + 1
```

This solution is SO good.
    """
    expected = "def solution(x):\n    return x + 1\n"
    assert find_code(completion) == expected


def test_find_code_with_markdown_block_and_text_before():
    completion = """
Let me solve this step by step:

```python
def solution(x):
    return x + 1
```
    """
    expected = "def solution(x):\n    return x + 1\n"
    assert find_code(completion) == expected


def test_find_code_without_markdown():
    completion = """
def solution(x):
    return x + 1
    """
    # When no markdown blocks found, returns original text with leading/trailing newlines
    expected = "\ndef solution(x):\n    return x + 1\n    "  # Match exact whitespace
    assert find_code(completion) == expected


def test_record_to_sample_basic():
    record = {
        "question": "Write a function that adds two numbers",
        "input_output": '{"inputs": [[1, 2], [3, 4]], "outputs": [3, 7]}',
        "problem_id": "123",
        "difficulty": "interview",
    }

    sample = record_to_sample(record)

    assert sample.target == [
        "assert solution([1, 2]) == 3",
        "assert solution([3, 4]) == 7",
    ]
    assert sample.id == "123"
    assert sample.metadata["problem_id"] == "123"
    assert sample.metadata["difficulty"] == "interview"
    assert (
        sample.metadata["test_list_str"]
        == "assert solution([1, 2]) == 3\nassert solution([3, 4]) == 7"
    )


def test_record_to_sample_multiple_tests():
    record = {
        "question": "Write a function that doubles a number",
        "input_output": '{"inputs": [1, 2, 3, 4], "outputs": [2, 4, 6, 8]}',
        "problem_id": "456",
        "difficulty": "interview",
    }

    sample = record_to_sample(record)

    # Check all test cases are in target
    assert len(sample.target) == 4
    assert sample.target == [
        "assert solution(1) == 2",
        "assert solution(2) == 4",
        "assert solution(3) == 6",
        "assert solution(4) == 8",
    ]

    # Check only first 3 tests are in prompt
    assert sample.metadata["test_list_str"] == (
        "assert solution(1) == 2\nassert solution(2) == 4\nassert solution(3) == 6"
    )


def test_record_to_sample_complex_types():
    record = {
        "question": "Write a function that processes lists",
        "input_output": '{"inputs": [[1, 2], ["a", "b"]], "outputs": [[2, 4], ["aa", "bb"]]}',
        "problem_id": "789",
        "difficulty": "interview",
    }

    sample = record_to_sample(record)

    assert sample.target == [
        "assert solution([1, 2]) == [2, 4]",
        "assert solution(['a', 'b']) == ['aa', 'bb']",
    ]


@pytest.mark.asyncio
async def test_verify_correct_code():
    """Test that verify() correctly scores a valid solution as correct"""
    # Create a mock task state with correct code
    state = TaskState(
        model="test_model",
        sample_id="test_sample",
        epoch=0,
        messages=[],
        input="test input",
        output=type(
            "MockOutput",
            (),
            {
                "completion": """```python
def solution(x):
    return x + 1
```"""
            },
        )(),
        metadata={},
    )

    # Create target with test cases
    target = Target(["assert solution(1) == 2", "assert solution(2) == 3"])

    # Get the scorer function
    scorer = verify()

    # Mock the sandbox execution to return success
    with patch("inspect_evals.apps.apps.sandbox") as mock_sandbox:
        mock_exec = AsyncMock(return_value=ExecResult(True, 0, "", ""))
        mock_sandbox.return_value.exec = mock_exec

        # Score the solution
        score = await scorer(state, target)

        assert score.value == CORRECT
        assert "All test cases passed" in score.explanation


@pytest.mark.asyncio
async def test_verify_incorrect_code():
    """Test that verify() correctly scores an invalid solution as incorrect"""
    # Create a mock task state with incorrect code
    state = TaskState(
        model="test_model",
        sample_id="test_sample",
        epoch=0,
        messages=[],
        input="test input",
        output=type(
            "MockOutput",
            (),
            {
                "completion": """```python
def solution(x):
    return x  # Wrong implementation
```"""
            },
        )(),
        metadata={},
    )

    # Create target with test cases
    target = Target(["assert solution(1) == 2", "assert solution(2) == 3"])

    # Get the scorer function
    scorer = verify()

    # Mock the sandbox execution to return failure
    with patch("inspect_evals.apps.apps.sandbox") as mock_sandbox:
        mock_exec = AsyncMock(
            return_value=ExecResult(False, 1, "", "AssertionError: assert 1 == 2")
        )
        mock_sandbox.return_value.exec = mock_exec

        # Score the solution
        score = await scorer(state, target)

        assert score.value == INCORRECT
        assert "Code did not pass all test cases" in score.explanation


@pytest.mark.asyncio
async def test_verify_timeout():
    """Test that verify() correctly handles timeouts by marking as incorrect"""
    # Create a mock task state
    state = TaskState(
        model="test_model",
        sample_id="test_sample",
        epoch=0,
        messages=[],
        input="test input",
        output=type(
            "MockOutput",
            (),
            {
                "completion": """```python
def solution(x):
    while True:  # Infinite loop
        pass
```"""
            },
        )(),
        metadata={},
    )

    # Create target with test cases
    target = Target(["assert solution(1) == 2"])

    # Get the scorer function
    scorer = verify()

    # Mock the sandbox execution to simulate timeout
    with patch("inspect_evals.apps.apps.sandbox") as mock_sandbox:
        mock_exec = AsyncMock(side_effect=TimeoutError())
        mock_sandbox.return_value.exec = mock_exec

        # Score the solution
        score = await scorer(state, target)

        assert score.value == INCORRECT
        assert "Verification timed out" in score.explanation
