from synthegrator.code_problems import CodeProblem
from synthegrator.synthdatasets import yield_problems_from_name, DatasetName
from synthegrator.synthdatasets.human_eval import yield_human_eval
import pytest
import tempfile
from synthegrator.synthdatasets import DatasetName


def test_yield_from_dataset_names():
    data_by_name = list(yield_problems_from_name(
        DatasetName.humaneval,
        max_problems=5
    ))
    data_by_func = list(yield_human_eval(max_problems=5))
    assert data_by_name[0].working_directory == data_by_func[0].working_directory
    assert are_probs_eq(data_by_name[0], data_by_func[0])


@pytest.mark.parametrize(
    "dataset",
    set(DatasetName) - {DatasetName.dypy_line_completion, DatasetName.mbpp_sanatized},
)
def test_load_each_dataset(dataset):
    print(dataset)
    data = list(yield_problems_from_name(
        dataset,
        max_problems=1
    ))
    assert len(data) == 1
    assert data[0].dataset_name == dataset.get_base_collection_name()

try:
    import joblib
    JOBLIB_AVAILABLE = True
except ImportError:
    JOBLIB_AVAILABLE = False

# Create a decorator to skip tests that require joblib
joblib_required = pytest.mark.skipif(
    not JOBLIB_AVAILABLE,
    reason="Test requires joblib to be installed"
)

# Example usage
@joblib_required
def test_joblib_caching_with_dataset_names():
    """Test that joblib's Memory correctly caches with DatasetName attributes."""
    # Setup temporary cache directory
    cache_dir = tempfile.mkdtemp()
    memory = joblib.Memory(cache_dir, verbose=0)

    # Counter for actual function calls
    call_count = 0

    @memory.cache
    def cached_function(dataset, value):
        nonlocal call_count
        call_count += 1
        return (dataset.name if hasattr(dataset, 'name') else dataset, value)

    # Test 1: First call should increment counter, second should use cache
    cached_function(DatasetName.humaneval, 42)
    assert call_count == 1, "First call should execute function"

    cached_function(DatasetName.humaneval, 42)
    assert call_count == 1, "Second call should use cache (not increment counter)"

    # Test 2: Different parameter should increment counter
    cached_function(DatasetName.humaneval, 43)
    assert call_count == 2, "Call with different value should execute function"

    # Test 3: Different dataset should increment counter
    cached_function(DatasetName.mbpp, 42)
    assert call_count == 3, "Call with different dataset should execute function"

    # Test 4: Dataset with parent relationship
    cached_function(DatasetName.humaneval_plus, 42)
    assert call_count == 4, "Call with parent dataset should execute function"

    # Call again with same parent dataset
    cached_function(DatasetName.humaneval_plus, 42)
    assert call_count == 4, "Repeated call with parent dataset should use cache"

    # Test 5: String representation equivalence
    call_count_before = call_count
    cached_function("humaneval", 42)
    assert call_count > call_count_before, "String representation should be treated differently"

    # Cleanup
    memory.clear()


def are_probs_eq(a: CodeProblem, b: CodeProblem):
    return bool(a == b)
