import random
from abc import ABC, abstractmethod

from synthegrator.code_problems import CodeProblem, TestCase


class PromptingTestCaseSelectionStrategy(ABC):
    """
    A strategy for selecting test cases to help the solver solve a problem.
    This is particularly necessary when doing programming-by-example, where
    the test cases are the main way the solver is specified the problem
    """

    @abstractmethod
    def select_for_problem(
        self,
        problem: CodeProblem,
        max_num_test_cases: int,
    ) -> list[TestCase]:
        raise NotImplementedError


class RandomVisibleTestCaseSelectionStrategy(
    PromptingTestCaseSelectionStrategy,
):
    def __init__(
        self,
        seed: int | None = None,
        cap_test_cases: int | None = None,
        randomize_order_even_when_less_than_cap: bool = False,
    ):
        self._seed = seed
        self._cap_test_cases = cap_test_cases
        self._randomize_order_even_when_less_than_cap = (
            randomize_order_even_when_less_than_cap
        )

    def select_for_problem(
        self,
        problem: CodeProblem,
        max_num_test_cases: int | None = None,
    ) -> list[TestCase]:
        if max_num_test_cases is None:
            max_num_test_cases = len(problem.test_cases)
        if self._cap_test_cases is not None:
            max_num_test_cases = min(max_num_test_cases, self._cap_test_cases)
        if max_num_test_cases <= 0:
            msg = "Max_num_test_cases must be greater than 0"
            raise ValueError(msg)

        available_test_cases = problem.get_non_hidden_test_cases()
        if (
            len(available_test_cases) <= max_num_test_cases
            and not self._randomize_order_even_when_less_than_cap
        ):
            return available_test_cases
        r = random.Random(
            f"{problem.problem_id}{problem.dataset_name}{self._seed}",
        )
        return r.choices(
            available_test_cases,
            k=min(max_num_test_cases, len(available_test_cases)),
        )
