from inspect_evals.usaco.usaco import load_test_cases, usaco

from .create_dataset import (
    IN_FILE_CONTENTS,
    OUT_FILE_CONTENTS,
    TEST_PROBLEM,
    TESTCASE_FILES,
    USACO_DATA_DIR,
    USACO_TESTCASE_DIR,
    disk_dataset,
    disk_testcases,
)

TEST_BASENAME = "usaco_test"


def verify_sample(sample, problem):
    """Ensure `sample` correctly represents `problem`."""
    # Verify id and input prompt
    assert sample.id == problem["cp_id"]
    assert problem["description"] in sample.input

    # Verify metadata
    assert sample.metadata["name"] == problem["name"]
    assert sample.metadata["problem_id"] == problem["problem_id"]
    assert sample.metadata["problem_level"] == problem["problem_level"]
    assert sample.metadata["runtime_limit_secs"] == problem["runtime_limit"]
    assert sample.metadata["memory_limit_mbs"] == problem["memory_limit"]
    assert sample.metadata["num_tests"] == problem["num_tests"]

    # Verify path to tests
    test_assets_path = str(
        sample.metadata["test_assets_path"]
    )  # convert Path object to str
    assert test_assets_path.startswith(str(USACO_DATA_DIR))
    assert sample.metadata["problem_id"] in test_assets_path


def test_uscao_task(test_basename=TEST_BASENAME):
    """Test task creation, configuration, and loaded dataset."""
    with disk_dataset(test_basename):
        task = usaco(json_basename=test_basename)  # calls read_dataset

    # Verify task configuration
    assert task.name == "usaco"
    assert task.dataset is not None
    assert len(task.dataset.samples) > 0

    # Check solver and scorer functions
    assert len(task.solver) == 1  # generate()
    assert len(task.scorer) == 1  # verify()

    verify_sample(task.dataset[0], TEST_PROBLEM)


def test_usaco_testcases():
    """Ensure the testcases are correctly loaded."""
    num_tests = len(TESTCASE_FILES)

    with disk_testcases():
        test_inputs, test_outputs = load_test_cases(USACO_TESTCASE_DIR, num_tests)

    for test_input, test_output in zip(test_inputs, test_outputs):
        assert type(test_input) is bytes  # since stdin uses bytes
        assert type(test_output) is str  # since stdout is str
        assert test_input == IN_FILE_CONTENTS
        assert test_output == OUT_FILE_CONTENTS.decode()
