from collections.abc import Iterator
from enum import Enum
from functools import partial
from itertools import islice

from tqdm import tqdm

from synthegrator.code_problem_builders import (
    make_instruction_only_problem,
    make_simple_method_completion_problem,
)
from synthegrator.code_problems import (
    CodeProblem,
    TestCase,
    TestCaseArbitraryMethod,
)
from synthegrator.lang_specs.lang_spec import LsFunctionDef
from synthegrator.lang_specs.lang_spec_python import PythonLangSpec
from synthegrator.problem_rendering import render_single_edit_codeproblem_autoregressive
from synthegrator.response_parser import set_base_indent
from synthegrator.solution_eval import evaluate_code_problem_execution
from synthegrator.synthdatasets import load_hf_dataset_with_backoff
from synthegrator.transformation_spec import (
    StsInsert,
)


class MbppSplit(Enum):
    TRAIN = "train"
    VALIDATION = "validation"
    TEST = "test"
    PROMPT = "prompt"
    ALL = "all"
    ALL_BUT_PROMPT = "all_but_prompt"
    """The mbpp dataset includes a split 10 problems that was their prompt for
    few shot solving. One might want to not do few shotting (or with alterative 
    leave-few-out versions), so we default to ALL including these. Use this version
    to exclude the 10 problems.
    """
    EVAL_PLUS = "eval_plus"
    """Load evalplus dataset with more test cases (evalplus.github.io). 
    Note though, this is only 374 problems, rather than the 974 
    found in actual MBPP"""


def yield_mbpp(
    split: MbppSplit = MbppSplit.ALL,
    reformat_whitespace: bool = True,
    sanitized_only: bool = False,
    only_with_passing_reference_solution: bool = False,
    format_as_method_completion_problem: bool = False,
    max_problems: int = None,
) -> Iterator[CodeProblem]:
    """
    :param split: which split to yield
    :param reformat_whitespace: whether to convert tabs to spaces and convert to unix line endings
    :param sanitized_only: whether to only yield sanitized examples. These are
        hand-verified by the authors.
    :param only_with_passing_reference_solution: About ~80/~970 of the problems don't
        seem to have a reference solution that does not pass the provided tests.
        You can set this to true to filter.
    :param format_as_method_completion_problem: Normally mbpp is file generation
        from instructions. Occasionally (about ~60 problems) this involves
        implemented a class or multiple methods. Setting this parameter will
        instead make the problem an insertion problem to fill in the contents
        of main method under test. Other context (like other classes, methods
        or imports) are provided as visible to the solver. This can make it
        easier to the solver as things like the imports can be a hint. It also
        makes it easier to get autoregressive language models to cooperate, as the
        method completion will always terminate at a block-end, whereas when
        just generating arbitrary completions from a model can end up with
        a lot of extraneous output that can be more difficult to filter out.
    :param max_problems: maximum number of problems to yield
    """
    # The all and the santized use different schemas
    dataset_all = _get_mbpp_huggingface(sanitized=False)
    """
    DatasetDict({
        test: Dataset({
            features: ['task_id', 'text', 'code', 'test_list', 'test_setup_code', 'challenge_test_list'],
            num_rows: 974
        })
    })
    """
    dataset_sanitized = _get_mbpp_huggingface(sanitized=True)
    """
    DatasetDict({
        test: Dataset({
            features: ['source_file', 'task_id', 'prompt', 'code', 'test_imports', 'test_list'],
            num_rows: 427
        })
    })
    """
    sanitized_key_to_instance = {
        ex["task_id"]: ex for ex in _iterate_for_split(dataset_sanitized, split)
    }

    exclude_task_ids = (
        _task_ids_where_reference_solution_fails()
        if only_with_passing_reference_solution
        else set()
    )

    produce_func = (
        _produce_code_problem_for_mbpp_method_completion
        if format_as_method_completion_problem
        else _produce_code_problem_for_mbpp
    )

    yielded_prob_count = 0
    for ex in _iterate_for_split(dataset_all, split):
        if max_problems is not None and yielded_prob_count >= max_problems:
            break
        santized_version = sanitized_key_to_instance.get(ex["task_id"], {})
        if sanitized_only and len(santized_version) == 0:
            continue
        if str(ex["task_id"]) in exclude_task_ids:
            continue
        try:
            prob = produce_func(
                prompt=santized_version.get("prompt", ex["text"]),
                code=ex["code"],
                challenge_test_list=ex["challenge_test_list"],
                task_id=ex["task_id"],
                test_list=santized_version.get("test_list", ex["test_list"]),
                test_setup_code=ex["test_setup_code"],
                test_imports=santized_version.get("test_imports", None),
                reformat_whitespace=reformat_whitespace,
                is_in_santized=len(santized_version) > 0,
            )
            yield prob
            yielded_prob_count += 1
        except DatasetParseException:
            # raise
            continue


def yield_mbpp_plus(
    reformat_whitespace: bool = True,
    format_as_method_completion_problem: bool = False,
    max_problems: int = None,
):
    """
    Load the MBPP+ dataset (evalplus.github.io). This dataset is intended
    to have more rigorous test cases than the original MBPP dataset. However,
    it has the disadvantage of being smaller (374 problems) compared to the
    974 problems in all splits of the original MBPP dataset.
    """
    ds = _get_mbpp_plus_huggingface()
    produce_func = (
        _produce_code_problem_for_mbpp_method_completion
        if format_as_method_completion_problem
        else _produce_code_problem_for_mbpp
    )
    assert list(ds.keys()) == ["test"]

    num_yielded = 0
    for ex in ds["test"]:
        if max_problems is not None and num_yielded >= max_problems:
            break
        try:
            prob = produce_func(
                prompt=ex["prompt"],
                code=ex["code"],
                challenge_test_list=[],
                task_id=ex["task_id"],
                test_preformated_src=ex["test"],
                test_list=ex["test_list"],
                test_setup_code="",
                test_imports=ex.get("test_imports", None),
                reformat_whitespace=reformat_whitespace,
                is_in_santized=False,
                is_plus=True,
            )
            assert prob.dataset_name == "mbpp_plus"
            yield prob
            num_yielded += 1
        except DatasetParseException as e:
            # raise
            continue


def _get_mbpp_huggingface(
    sanitized: bool,
):
    if not sanitized:
        return load_hf_dataset_with_backoff("mbpp", "full")
    return load_hf_dataset_with_backoff(
        "mbpp",
        "sanitized",
    )


def _get_mbpp_plus_huggingface():
    data = load_hf_dataset_with_backoff("evalplus/mbppplus")
    return data


def _iterate_for_split(dataset, split: MbppSplit):
    all_main = split in {MbppSplit.ALL, MbppSplit.ALL_BUT_PROMPT}
    if split == MbppSplit.TRAIN or all_main:
        yield from dataset["train"]
    if split == MbppSplit.VALIDATION or all_main:
        yield from dataset["validation"]
    if split == MbppSplit.TEST or all_main:
        yield from dataset["test"]
    if split in {MbppSplit.PROMPT, MbppSplit.ALL}:
        yield from dataset["prompt"]


def _task_ids_where_reference_solution_fails():
    return {
        "602",
        "617",
        "620",
        "626",
        "627",
        "631",
        "633",
        "643",
        "721",
        "734",
        "737",
        "743",
        "746",
        "747",
        "751",
        "755",
        "756",
        "757",
        "760",
        "763",
        "767",
        "770",
        "777",
        "781",
        "782",
        "787",
        "793",
        "794",
        "797",
        "799",
        "801",
        "802",
        "803",
        "804",
        "521",
        "576",
        "586",
        "597",
        "16",
        "19",
        "46",
        "51",
        "57",
        "61",
        "70",
        "75",
        "92",
        "101",
        "117",
        "119",
        "125",
        "130",
        "137",
        "145",
        "159",
        "160",
        "164",
        "166",
        "167",
        "172",
        "245",
        "264",
        "270",
        "285",
        "296",
        "327",
        "406",
        "433",
        "434",
        "436",
        "438",
        "452",
        "454",
        "468",
        "471",
    }


_mbpp_insert_node = StsInsert(stop_at_block_end=False)


def _produce_code_problem_for_mbpp(
    prompt: str,
    code: str,  # solution code
    challenge_test_list: list[str],
    task_id: int,
    test_list: list[str],
    test_setup_code: str,
    reformat_whitespace: bool,
    test_imports: list[str],
    is_in_santized: bool,
    test_preformated_src: str = None,
    is_plus: bool = False,
) -> CodeProblem:
    # Copy prompt from https://github.com/google-research/google-research/tree/master/mbpp
    # You are an expert Python programmer, and here is your task: {prompt} Your code should pass these tests:\n\n{tests}\n[BEGIN]\n{code}\n[DONE
    test_prompt = "\n".join(test_list)
    prompt = (
        f"You are an expert Python programmer, and here is your task: {prompt} "
        "Your code should pass these tests:\n\n"
        f"{test_prompt}\n\n"
    )
    if reformat_whitespace:
        prompt = _reformat_whitespace(prompt)

    if not test_preformated_src:
        test_cases = _make_mbpp_test_cases(
            test_list,
            challenge_test_list,
            test_imports,
            test_setup_code,
            reformat_whitespace,
        )
    else:
        test_cases = _make_mbpp_plus_test_cases(
            test_list,
            challenge_test_list,
            test_imports,
            test_setup_code,
            reformat_whitespace,
            test_reformated_src=test_preformated_src,
        )

    return make_instruction_only_problem(
        instructions=prompt,
        insert_node=_mbpp_insert_node,
        test_cases=test_cases,
        dataset_name="mbpp" if not is_plus else "mbpp_plus",
        problem_id=str(task_id),
        known_solutions=[code],
    )


def _produce_code_problem_for_mbpp_method_completion(
    prompt: str,
    code: str,  # solution code
    challenge_test_list: list[str],
    task_id: int,
    test_list: list[str],
    test_setup_code: str,
    reformat_whitespace: bool,
    test_imports: list[str],
    is_in_santized: bool,
    test_preformated_src: str = None,
    is_plus: bool = False,
) -> CodeProblem:
    test_prompt = "\n".join(test_list)
    reformat = _reformat_whitespace if reformat_whitespace else lambda x: x
    prompt = reformat(f"{prompt}\n\nExamples:\n{test_prompt}")
    ref_func = _get_reference_method_under_test(reformat(code), test_list)
    if not test_preformated_src:
        test_cases = _make_mbpp_test_cases(
            test_list,
            challenge_test_list,
            test_imports,
            test_setup_code,
            reformat_whitespace,
        )
    else:
        test_cases = _make_mbpp_plus_test_cases(
            test_list,
            challenge_test_list,
            test_imports,
            test_setup_code,
            reformat_whitespace,
            test_reformated_src=test_preformated_src,
        )

    return make_simple_method_completion_problem(
        prompt=prompt,
        signature=reformat(ref_func.get_declaration_str()),
        model_context=reformat(ref_func.get_preceding_src_code()).lstrip("\n"),
        known_solutions=[
            set_base_indent(reformat(ref_func.get_body_src()).lstrip("\n"), 4),
        ],
        after_method_src=ref_func.get_src_code_after().strip(),
        test_cases=test_cases,
        dataset_name="mbpp" if not is_plus else "mbpp_plus",
        problem_id=str(task_id) + "func",
    )


class DatasetParseException(Exception):
    pass


def _get_reference_method_under_test(reference_code, test_list) -> LsFunctionDef:
    method_name = _get_mbpp_method_name_under_test(test_list)
    if not method_name:
        msg = (
            f"Expected to find method name.\nReference code:\n{reference_code}\nTest"
            f" list:\n{test_list}"
        )
        raise DatasetParseException(
            msg,
        )
    funcs = list(PythonLangSpec().find_functions(reference_code))
    for func in funcs:
        if func.get_function_name() == method_name:
            return func
    msg = f"Unable to find the mbpp method under test from the test list:\n{test_list}"
    raise DatasetParseException(
        msg,
    )


def _get_mbpp_method_name_under_test(test_list: list[str]) -> str:
    def _guess_from_assert_statement(statement: str):
        if statement.startswith("assert ("):
            # A weird case where just has an unnessary paren (ex, 769)
            statement = statement[len("assert (") :]
        else:
            statement = statement[len("assert ") :]
        # just guess that it is before the first open paren. This could
        #   fail if wrapped in some other call, but will see if the same for all the tests
        return statement.split("(", 1)[0].strip()

    guesses = {_guess_from_assert_statement(stmt) for stmt in test_list}
    if len(guesses) != 1:
        msg = (
            "Unable to guess the mbpp method under test from the test"
            f" list:\n{test_list}"
        )
        raise DatasetParseException(
            msg,
        )
    value = guesses.pop()
    common_builtins = {
        "max",
        "min",
        "list",
        "tuple",
        "str",
        "dict",
        "set",
        "len",
        "sum",
        "sorted",
        "reversed",
        "enumerate",
    }
    if value in common_builtins:
        msg = (
            "The mbpp method under test is a common builtin. This is probably a bad"
            " parse."
        )
        raise DatasetParseException(
            msg,
        )
    return value


def _make_mbpp_test_cases(
    test_list,
    challenge_test_list,
    test_imports,
    test_setup_code,
    reformat_whitespace,
):
    def make_tc(tc_list, is_hidden):
        return _parse_test_cases(
            test_cases=tc_list,
            test_case_imports=test_imports,
            test_case_setup=test_setup_code,
            is_hidden_test=is_hidden,
            reformat_whitespace=reformat_whitespace,
        )

    return make_tc(test_list, False) + make_tc(challenge_test_list, True)


def _make_mbpp_plus_test_cases(
    test_list,
    challenge_test_list,
    test_imports,
    test_setup_code,
    reformat_whitespace,
    test_reformated_src,
):
    base_tests = _make_mbpp_test_cases(
        test_list, [], test_imports, test_setup_code, reformat_whitespace
    )
    challenge = [
        TestCaseArbitraryMethod(
            is_hidden_test=True,
            fail_message="Fail challenge mbpp test",
            method_body=test_reformated_src,
            lang_spec=PythonLangSpec(),
        )
    ]
    return base_tests + challenge


def _reformat_whitespace(code: str) -> str:
    """Converts tabs to spaces and changes to unix line endings"""
    return code.replace("\t", "    ").replace("\r\n", "\n")


def _parse_test_cases(
    test_cases: list[str],
    test_case_imports: list[str],
    test_case_setup: str,
    is_hidden_test: bool,
    reformat_whitespace: bool,
) -> list[TestCase]:
    out = []
    for tc in test_cases:
        tc_imports = "\n".join(test_case_imports or [])
        body = f"{tc_imports}\n{test_case_setup}\n{tc}"
        if reformat_whitespace:
            body = _reformat_whitespace(body)
        body = body.lstrip()
        out.append(
            TestCaseArbitraryMethod(
                is_hidden_test=is_hidden_test,
                fail_message=(
                    f"Fail {'hidden' if is_hidden_test else 'unhidden'} test case: {tc}"
                ),
                method_body=body,
                lang_spec=PythonLangSpec(),
            ),
        )
    return out


def print_problems_with_failing_reference_solutions():
    bad_probs = []
    for problem in tqdm(
        list(
            yield_mbpp(
                split=MbppSplit.ALL,
                sanitized_only=False,
                only_with_passing_reference_solution=False,
            ),
        ),
    ):
        if len(problem.known_solutions) == 0:
            bad_probs.append(problem)
        result = evaluate_code_problem_execution(
            problem,
            problem.known_solutions[0],
            do_not_execute_if_syntax_fail=True,
        )
        if not result.all_passed():
            bad_probs.append(problem)
    print("Problems with failing reference solutions:")
    prob_ids = [p.problem_id for p in bad_probs]
    print(repr(prob_ids))
    assert set(prob_ids) == set(_task_ids_where_reference_solution_fails())
    print("Matches current list")


def _eval_failing():
    fails = [
        prob
        for prob in yield_mbpp(
            split=MbppSplit.ALL,
            sanitized_only=False,
            only_with_passing_reference_solution=False,
        )
        if prob.problem_id in _task_ids_where_reference_solution_fails()
    ]
    for problem in tqdm(fails):
        r = evaluate_code_problem_execution(
            problem,
            problem.known_solutions[0],
            do_not_execute_if_syntax_fail=True,
        )
        print("Instructions:")
        print(problem.instructions)
        print("Render:")
        print(render_single_edit_codeproblem_autoregressive(problem).text)
        assert not r.all_passed()


def _measure_method_completion_probs():
    """
    A quick script to try to estimate how many of them seem to be method completion
    problems
    """
    from datasets import load_dataset

    dataset_all = load_dataset("mbpp")
    lang_spec = PythonLangSpec()
    lang = lang_spec.get_tree_sitter_lang()
    num_single_func = 0
    num_has_imports_single_func = 0
    num_total = 0
    num_has_imports = 0
    for ex in _iterate_for_split(dataset_all, MbppSplit.ALL):
        tree = lang_spec.get_tree_sitter_tree(ex["code"])
        # find all the function definitions
        query = lang.query("(function_definition name: (identifier) @function)")
        captures = query.captures(tree.root_node)
        len(list(captures))
        def_count = ex["code"].count("def ")
        class_count = ex["code"].count("class ")
        import_count = ex["code"].count("import ")
        has_imports = import_count > 0
        if def_count == 1 and class_count == 0:
            # for line in ex["code"].splitlines():
            #    if line.startswith("def "):
            #        assert line.strip().endswith(":")
            num_single_func += 1
            if has_imports:
                num_has_imports_single_func += 1
        num_total += 1
        if has_imports:
            num_has_imports += 1
    print("num_total", num_total)
    print("num_single_func", num_single_func)
    print("not single func", num_total - num_single_func)
    print("num_has_imports_single_func", num_has_imports_single_func)
    print("num_has_imports", num_has_imports)
    print(len(_task_ids_where_reference_solution_fails()))


if __name__ == "__main__":
    # print_problems_with_failing_reference_solutions()
    _measure_method_completion_probs()
