import base64

import pytest
import pytest_asyncio  # noqa: F401

from inspect_evals.cyberseceval_2.vulnerability_exploit.challenges.memory_corruption.examples import (
    CAdvancedMemoryCorruptionGenerator,
    CodeRunner,
)
from inspect_evals.cyberseceval_2.vulnerability_exploit.challenges.memory_corruption.generator import (
    MemoryCorruptionTestPrompt,
    generate_advanced_memory_corruption_test_cases,
)

NUM_TEST_CASES = 10


def memory_corruption_test_cases():
    gen = CAdvancedMemoryCorruptionGenerator()
    return gen.generate_static_memory_corruption_test_cases(NUM_TEST_CASES)


def adv_memory_corruption_test_cases():
    return generate_advanced_memory_corruption_test_cases(NUM_TEST_CASES)


def sample_answer(a) -> bytes:
    a = base64.b64decode(a)
    if a[0] == ord("b"):
        a = a[2:-1]
    return a.decode("unicode_escape").encode("ISO 8859-1")


@pytest.mark.parametrize(
    "test_case", memory_corruption_test_cases(), ids=lambda x: x["name"]
)
def test_memory_corruption(test_case: MemoryCorruptionTestPrompt):
    # FIXME: `example_heap_buffer_overflow` and `example_format_string` are consistently failing (on MacOS 14.5)
    assert len(test_case["code"]) > 20
    assert len(test_case["mutated_prompt"]) > 120
    assert len(test_case["answer"]) > 0

    if test_case["name"] in (
        "example_hash_crc32",
        "example_hash_crc32_hex",
        "example_iterator_invalidation",
    ):
        # These tests fail on CI but pass locally
        # https://github.com/UKGovernmentBEIS/inspect_evals/issues/168
        pytest.xfail("Issue #168: Fails on CI")

    runner = CodeRunner(source_code=test_case["code"], name=test_case["name"])

    # Test with sample answer
    exit_code, output = runner.check_output(sample_answer(test_case["answer"]))
    assert exit_code != 0, (
        f"Expected code to crash with sample answer but got:\n{output}"
    )

    # Test with empty input
    exit_code, output = runner.check_output(b"")
    assert exit_code == 0, f"Code should not crash with empty input but got:\n{output}"


@pytest.mark.skip("Advanced memory corruption tests are flakey - see Issue #168")
@pytest.mark.parametrize(
    "test_case", adv_memory_corruption_test_cases(), ids=lambda x: x["name"]
)
def test_adv_memory_corruption(test_case: MemoryCorruptionTestPrompt):
    # FIXME: These tests are flakey, and will sometimes pass when retried (on MacOS 14.5)
    # All of them have succeeded at least once.
    assert len(test_case["code"]) > 20
    assert len(test_case["mutated_prompt"]) > 120
    assert len(test_case["answer"]) > 0

    runner = CodeRunner(source_code=test_case["code"], name=test_case["name"])

    # Test with sample answer
    exit_code, output = runner.check_output(sample_answer(test_case["answer"]))
    assert exit_code != 0, (
        f"Expected code to crash with sample answer but got:\n{output}"
    )

    # Test with empty input
    exit_code, output = runner.check_output(b"")
    assert exit_code == 0, f"Code should not crash with empty input but got:\n{output}"
