import pytest
from inspect_ai.scorer import CORRECT, INCORRECT
from inspect_ai.util import ExecResult

from inspect_evals.usaco.usaco import run_test, run_tests

IN_FILE_CONTENTS = b"1\n"
OUT_FILE_CONTENTS = b"2\n"  # e.g. increment the integer in IN_FILE


@pytest.mark.asyncio
async def test_usaco_run_single_tests():
    code = 'print("Hi!")'
    test_input = b"1\n"
    test_output = "2\n"

    # PASS: stdout matches test_output
    result = ExecResult(success=True, returncode=0, stdout=test_output, stderr="")
    passed, explanation = await run_test(code, test_input, test_output, result=result)
    assert explanation != ""
    assert passed is True

    # FAIL: stdout doesn't match test_output
    result = ExecResult(success=True, returncode=0, stdout="", stderr="")
    passed, explanation = await run_test(code, test_input, test_output, result=result)
    assert explanation != ""
    assert passed is False

    # FAIL: stdout matches, but returncode not 0
    result = ExecResult(
        success=False, returncode=1, stdout=test_output, stderr="exit status 1"
    )
    passed, explanation = await run_test(
        code, test_input, test_output, result=result, verbose=False
    )
    assert explanation == ""
    assert passed is False


@pytest.mark.asyncio
async def test_usaco_run_tests():
    test_output = OUT_FILE_CONTENTS
    results = [
        ExecResult(success=True, returncode=0, stdout=test_output, stderr=""),  # True
        ExecResult(success=True, returncode=0, stdout=test_output, stderr=""),  # True
        ExecResult(
            success=True, returncode=0, stdout="", stderr=""
        ),  # False, so should early stop here
        ExecResult(success=True, returncode=0, stdout=test_output, stderr=""),  # True
        ExecResult(
            success=False, returncode=1, stdout=test_output, stderr="exit status 1"
        ),  # False
    ]

    code = 'print("2")'
    test_inputs = [IN_FILE_CONTENTS] * len(results)
    test_outputs = [test_output] * len(results)

    # TEST SUCCESS
    test_result, n_passed, explanation = await run_tests(
        code, test_inputs[:2], test_outputs[:2], results=results[:2]
    )
    assert n_passed == 2
    assert explanation != ""
    assert test_result is CORRECT

    # TEST FAILURE -- verify early stopping
    test_result, n_passed, explanation = await run_tests(
        code, test_inputs, test_outputs, results=results
    )
    assert n_passed == 2
    assert explanation != ""
    assert test_result is INCORRECT

    # TEST FAILURE -- same but verbose is False, so no explanation
    test_result, n_passed, explanation = await run_tests(
        code, test_inputs, test_outputs, results=results, verbose=False
    )
    assert n_passed == 2
    assert explanation == ""
    assert test_result is INCORRECT
