import dataclasses
import difflib
from itertools import islice
from pprint import pprint
from unittest.mock import MagicMock

import pytest

from synthegrator.code_problems import CodeProblem
from synthegrator.code_solver import DummyCodeSolverAutoRegressive
from synthegrator.lang_specs.lang_spec_python import PythonLangSpec
from synthegrator.problem_rendering import (
    PromptRendererSingleEditGeneric,
    render_single_edit_codeproblem_autoregressive,
)
from synthegrator.solution_eval import apply_solution, evaluate_code_problem_execution
from synthegrator.synthdatasets.mbpp import (
    MbppSplit,
    _get_mbpp_huggingface,
    _get_mbpp_plus_huggingface,
    _iterate_for_split,
    yield_mbpp,
    yield_mbpp_plus,
)
from synthegrator.transformation_spec import SolveStep
from synthegrator.util import pretty_print_python_code


def test_get_single_instance():
    val = next(yield_mbpp())
    assert val.problem_id == "601"

    assert (
        val.instructions
        == "You are an expert Python programmer, and here is your task: Write a"
        " function to find the longest chain which can be formed from the given set"
        " of pairs. Your code should pass these tests:\n\nassert"
        " max_chain_length([Pair(5, 24), Pair(15, 25),Pair(27, 40), Pair(50, 60)],"
        " 4) == 3\nassert max_chain_length([Pair(1, 2), Pair(3, 4),Pair(5, 6),"
        " Pair(7, 8)], 4) == 4\nassert max_chain_length([Pair(19, 10), Pair(11,"
        " 12),Pair(13, 14), Pair(15, 16), Pair(31, 54)], 5) == 5\n\n"
    )
    assert val.dataset_name == "mbpp"


def test_get_single_instance_method_completion():
    val = next(yield_mbpp(format_as_method_completion_problem=True))
    assert val.problem_id == "601func"
    assert val.dataset_name == "mbpp"
    render = render_single_edit_codeproblem_autoregressive(val)
    pretty_print_python_code(render.prompt.text)
    # fmt: off
    assert render.prompt.text == (
        'class Pair(object): \n'
        '    def __init__(self, a, b): \n'
        '        self.a = a \n'
        '        self.b = b \n'
        'def max_chain_length(arr, n): \n'
        '    """\n'
        '    Write a function to find the longest chain which can be formed from the given set of pairs.\n'
        '    \n'
        '    Examples:\n'
        '    assert max_chain_length([Pair(5, 24), Pair(15, 25),Pair(27, 40), Pair(50, 60)], 4) == 3\n'
        '    assert max_chain_length([Pair(1, 2), Pair(3, 4),Pair(5, 6), Pair(7, 8)], 4) == 4\n'
        '    assert max_chain_length([Pair(19, 10), Pair(11, 12),Pair(13, 14), Pair(15, 16), Pair(31, 54)], 5) == 5\n'
        '    """\n'
    )
    sol = apply_solution(val.known_solutions[0]).get_only_file().content_str
    assert sol == (
        'class Pair(object): \n'
        '    def __init__(self, a, b): \n'
        '        self.a = a \n'
        '        self.b = b \n'
        'def max_chain_length(arr, n): \n'
        '    """\n'
        '    Write a function to find the longest chain which can be formed from the given set of pairs.\n'
        '    \n'
        '    Examples:\n'
        '    assert max_chain_length([Pair(5, 24), Pair(15, 25),Pair(27, 40), Pair(50, 60)], 4) == 3\n'
        '    assert max_chain_length([Pair(1, 2), Pair(3, 4),Pair(5, 6), Pair(7, 8)], 4) == 4\n'
        '    assert max_chain_length([Pair(19, 10), Pair(11, 12),Pair(13, 14), Pair(15, 16), Pair(31, 54)], 5) == 5\n'
        '    """\n'
        '    max = 0\n'
        '    mcl = [1 for i in range(n)] \n'
        '    for i in range(1, n): \n'
        '        for j in range(0, i): \n'
        '            if (arr[i].a > arr[j].b and\n'
        '                mcl[i] < mcl[j] + 1): \n'
        '                mcl[i] = mcl[j] + 1\n'
        '    for i in range(n): \n'
        '        if (max < mcl[i]): \n'
        '            max = mcl[i] \n'
        '    return max\n'
    )
    # fmt: on


@pytest.mark.parametrize(
    "format_as_method_completion_problem",
    [False, True],
    ids=["instruction", "method_completion"],
)
def test_get_all(format_as_method_completion_problem):
    assert len(list(yield_mbpp(MbppSplit.TRAIN))) == 374
    assert len(list(yield_mbpp(MbppSplit.VALIDATION))) == 90
    assert len(list(yield_mbpp())) == 974
    assert len(list(yield_mbpp(sanitized_only=True))) == 427


def test_apply_prime_num():
    for prob in islice(
        yield_mbpp(
            MbppSplit.ALL,
            only_with_passing_reference_solution=True,
            format_as_method_completion_problem=True,
        ),
        3,
        4,
    ):
        assert len(prob.known_solutions) == 1
        assert prob.problem_id == "605func"  # The method prime_num problem which has issues
        apply_solution(prob.known_solutions[0])


# def test_apply_prob911():
#    # This problem has a `assert int(fuc())` style
#    for prob in islice(yield_mbpp(
#        MbppSplit.ALL,
#    ), 911 - 601, 911 - 601 + 1):


@pytest.mark.parametrize(
    "format_as_method_completion_problem",
    [False, True],
    ids=["instruction", "method_completion"],
)
def test_apply_all_solutions(format_as_method_completion_problem):
    for prob in yield_mbpp(
        MbppSplit.ALL,
        only_with_passing_reference_solution=True,
        format_as_method_completion_problem=format_as_method_completion_problem,
    ):
        assert len(prob.known_solutions) == 1
        print(prob.problem_id)
        apply_solution(prob.known_solutions[0])


def test_whitespace_on_prob():
    for prob in yield_mbpp(format_as_method_completion_problem=True):
        if prob.problem_id == "226func":
            break
    else:
        msg = "Problem not found"
        raise AssertionError(msg)
    renderer = PromptRendererSingleEditGeneric()
    lm = MagicMock()
    lm.is_chat_model = False
    render = renderer.render(prob, lm)
    # fmt: off
    assert render.prompt.text == (
        '''def odd_values_string(str):
    """
    Write a python function to remove the characters which have odd index values of a given string.
    
    Examples:
    assert odd_values_string('abcdef') == 'ace'
    assert odd_values_string('python') == 'pto'
    assert odd_values_string('data') == 'dt'
    assert odd_values_string('lambs') == 'lms'
    """
'''
    )
    # fmt: on


@pytest.mark.parametrize(
    "format_as_method_completion_problem",
    [False, True],
    ids=["instruction", "method_completion"],
)
def test_dummy_solve(format_as_method_completion_problem):
    dummy_solve = DummyCodeSolverAutoRegressive(MagicMock())
    failing_probs = []
    real_exs = _get_mbpp_huggingface(sanitized=False)
    real_exs = {
        str(ex["task_id"]): ex for ex in _iterate_for_split(real_exs, MbppSplit.ALL)
    }
    for prob in islice(
        yield_mbpp(
            MbppSplit.ALL,
            only_with_passing_reference_solution=True,
            format_as_method_completion_problem=format_as_method_completion_problem,
        ),
        None,
        None,
        100,
    ):
        solution = dummy_solve.solve(prob, PythonLangSpec())
        result = evaluate_code_problem_execution(prob, solution)
        if not result.all_passed():
            print("problem id")
            print(prob.problem_id)
            print("--- Application:")
            pretty_print_python_code(
                apply_solution(solution).get_only_file().content_str,
            )
            print(result)
            print(prob)
            print("Tc")
            print(result)
            print("Tc")
            for tc in real_exs[prob.problem_id]["test_list"]:
                pretty_print_python_code(tc)
            print("Instructions:")
            print(prob.instructions)
            print("-- Real")
            pretty_print_python_code(real_exs[prob.problem_id]["code"])
            print("-- Render:")
            pretty_print_python_code(
                render_single_edit_codeproblem_autoregressive(prob).text,
            )
            failing_probs.append(prob)
            raise AssertionError
    print(failing_probs)
    print("failing ids", [p.problem_id for p in failing_probs])
    assert len(failing_probs) == 0


@pytest.mark.parametrize(
    "format_as_method_completion_problem",
    # [False, True],
    # ids=["instruction", "method_completion"],
    [False],
    ids=["instruction"],
)
def test_mbpp_plus(format_as_method_completion_problem):
    dummy_solve = DummyCodeSolverAutoRegressive(MagicMock())
    failing_probs = []
    real_exs = _get_mbpp_plus_huggingface()
    real_exs = {str(ex["task_id"]): ex for ex in real_exs["test"]}
    probs = list(
        yield_mbpp_plus(
            format_as_method_completion_problem=format_as_method_completion_problem,
        )
    )
    assert len(probs) == 378
    for prob in islice(
        probs,
        None,
        None,
        100,
    ):
        solution = dummy_solve.solve(prob, PythonLangSpec())
        result = evaluate_code_problem_execution(prob, solution)
        if not result.all_passed():
            print("problem id")
            print(prob.problem_id)
            print("--- Application:")
            pretty_print_python_code(
                apply_solution(solution).get_only_file().content_str,
            )
            print(result)
            print(prob)
            print("Tc")
            print(result)
            print("Tc")
            for tc in real_exs[prob.problem_id]["test_list"]:
                pretty_print_python_code(tc)
            print("Instructions:")
            print(prob.instructions)
            print("-- Real")
            pretty_print_python_code(real_exs[prob.problem_id]["code"])
            print("-- Render:")
            pretty_print_python_code(
                render_single_edit_codeproblem_autoregressive(prob).text,
            )
            failing_probs.append(prob)
            raise AssertionError
    print(failing_probs)
    print("failing ids", [p.problem_id for p in failing_probs])
    assert len(failing_probs) == 0


def test_a_solve_that_passes_mbpp_but_not_plus():
    first_plus_prob = next(
        yield_mbpp_plus(
            format_as_method_completion_problem=True,
        )
    )
    probs = list(
        yield_mbpp(
            split=MbppSplit.ALL,
            format_as_method_completion_problem=True,
        )
    )
    mbpp_instruct_to_p: dict[str, CodeProblem] = {
        p.working_directory.files.get_only_file().content_str: p for p in probs
    }
    assert len(probs) == len(mbpp_instruct_to_p)
    pprint("all problems")
    print("problem id", first_plus_prob.instructions)
    content = first_plus_prob.working_directory.files.get_only_file().content_str
    print("Content")
    print("{" + content + "}")
    for key, prob in mbpp_instruct_to_p.items():
        if "is_not_prime" in key:
            print("FOUND")
            print("{" + key + "}")
            for tc in prob.test_cases:
                print("TC", tc.method_body)
            """
            TC assert is_not_prime(2) == False
            TC assert is_not_prime(10) == True
            TC assert is_not_prime(35) == True
            TC assert is_not_prime(37) == False
            """
            try:
                assert key == content
            except AssertionError:
                diff = difflib.unified_diff(
                    key.splitlines(keepends=True),
                    content.splitlines(keepends=True),
                    fromfile="key",
                    tofile="content",
                    lineterm="",
                )
                diff_text = "".join(diff)
                stripped_match = key.strip() == content.strip()
                raise AssertionError(
                    f"Strings are not equal. Stripped match: {stripped_match}. Diff:\n{diff_text}"
                ) from None
    normal_version = mbpp_instruct_to_p[content]
    print(normal_version)
    print("normal version")
    print(normal_version.instructions)
    print(print(normal_version))
    # fmt: off
    stupid_version = (
"""
    return {
        2: False,
        10: True,
        35: True,
        37: False,
    }.get(n, True)
"""
    )
    # fmt: on
    normal_sol = normal_version.known_solutions[0]
    stupid_solve = dataclasses.replace(
        normal_sol,
        solve_steps=[
            SolveStep(
                path=normal_sol.solve_steps[0].path,
                mark_id=normal_sol.solve_steps[0].mark_id,
                value=stupid_version,
            )
        ],
    )
    stupid_eval_normal = evaluate_code_problem_execution(normal_version, stupid_solve)
    assert stupid_eval_normal.all_passed()
    new_solve = normal_version.known_solutions[0]
    stupid_solve = dataclasses.replace(
        new_solve,
        solve_steps=[
            SolveStep(
                path=new_solve.solve_steps[0].path,
                mark_id=new_solve.solve_steps[0].mark_id,
                value=stupid_version,
            )
        ],
    )
    stupid_eval_new = evaluate_code_problem_execution(first_plus_prob, stupid_solve)
    assert not stupid_eval_new.all_passed()
