from unittest.mock import MagicMock

import pytest

from synthegrator.code_problem_builders import (
    make_simple_method_completion_problem,
    make_solve_for_single_edit,
)
from synthegrator.code_problems import (
    CodeProblem,
    CodeSolution,
    DiscoveredTestsuite,
    TestCaseMethodCallIsEq,
)
from synthegrator.environments import ProjectWorkingDirectory
from synthegrator.lang_specs.lang_spec_python import PythonLangSpec
from synthegrator.memory_fs import ProjectDir
from synthegrator.problem_rendering import render_single_edit_codeproblem_autoregressive
from synthegrator.solution_eval import (
    TestRunResultList,
    apply_solution,
    apply_steps_to_markup,
    create_exact_match_metric,
    evaluate_code_problem_execution,
    parse_junit_test_cases,
)
from synthegrator.tests.non_dataset_tests.test_value_question_render import (
    basic_classification_question_problem,
)
from synthegrator.transformation_spec import (
    MarkElement,
    MarkText,
    SolveStep,
    StsInsert,
    StsInvisible,
    StsNormalText,
    StsSelectorCharSpan,
    StsSelectorLineSpan,
    StsSpecStatement,
    TransformationSpec,
)
from synthegrator.util import get_only_item


def test_evaluate_code_problem():
    problem = make_simple_method_completion_problem(
        signature="def foo(x) -> int:",
        prompt="Write a function that doubles its input.",
        test_cases=[
            TestCaseMethodCallIsEq(
                test_input=(2,),
                output=4,
                method_name="foo",
                is_hidden_test=False,
                fail_message="",
            ),
            TestCaseMethodCallIsEq(
                test_input=(3,),
                output=6,
                method_name="foo",
                is_hidden_test=False,
                fail_message="",
            ),
        ],
    )
    solve = make_solve_for_single_edit(problem, "    return x * 2")
    result_list = evaluate_code_problem_execution(
        problem,
        solve,
        user_prompted_guard=False,
    )
    assert result_list[0].success
    assert result_list[1].success
    assert result_list.all_passed()
    solve_eval = problem.preferred_solution_evaluator.evaluate(solve)
    assert solve_eval.main_metric.is_success
    assert solve_eval.main_metric.float_val == 1.0
    assert solve_eval.main_metric.label == TestRunResultList.METRIC_NAME


def test_evaluate_code_problem_wrong():
    problem = make_simple_method_completion_problem(
        signature="def foo(x) -> int:",
        prompt="Write a function that doubles its input.",
        test_cases=[
            TestCaseMethodCallIsEq(
                test_input=(2,),
                output=4,
                method_name="foo",
                is_hidden_test=False,
                fail_message="",
            ),
            TestCaseMethodCallIsEq(
                test_input=(3,),
                output=6,
                method_name="foo",
                is_hidden_test=False,
                fail_message="",
            ),
        ],
    )
    solve = make_solve_for_single_edit(problem, "    return 4")
    result_list = evaluate_code_problem_execution(
        problem,
        solve,
        user_prompted_guard=False,
    )
    assert result_list[0].success
    assert not result_list[1].success
    assert not result_list.all_passed()
    solve_eval = problem.preferred_solution_evaluator.evaluate(solve)
    assert not solve_eval.main_metric.is_success
    assert solve_eval.main_metric.float_val == 0.5
    assert solve_eval.main_metric.label == TestRunResultList.METRIC_NAME


def test_apply_steps_to_markup():
    node = MarkElement(
        StsNormalText(),
        [
            MarkText("foo", 0, 3),
            MarkElement(StsInvisible(), [MarkText("bar", 3, 6)], 3, 6),
            MarkText("baz", 6, 9),
            MarkElement(StsInsert(), [], 9, 9),
        ],
        0,
        9,
    ).freeze()
    insert_node = node.children[3]
    assert isinstance(insert_node.verb, StsInsert)
    val = apply_steps_to_markup(node, [])
    assert val == "foobarbaz"
    val = apply_steps_to_markup(node, [SolveStep("foo.py", insert_node.mark_id, "qux")])
    assert val == "foobarbazqux"
    val = apply_steps_to_markup(node, [SolveStep("foo.py", "wrong", "qux")])
    assert val == "foobarbaz"
    val = apply_steps_to_markup(
        node,
        [SolveStep("foo.py", insert_node.mark_id, "qux")],
        still_hide_invisible=True,
    )
    assert val == "foobazqux"


def test_apply_steps_to_markup2():
    node = MarkElement(
        StsNormalText(),
        [
            MarkText("foo", 0, 3),
            MarkElement(StsInvisible(), [MarkText("bar", 3, 6)], 3, 6),
            MarkElement(StsInsert(), [], 6, 6),
            MarkText("baz", 6, 9),
        ],
        0,
        9,
    ).freeze()
    insert_node = node.children[2]
    assert isinstance(insert_node.verb, StsInsert)
    val = apply_steps_to_markup(node, [])
    assert val == "foobarbaz"
    val = apply_steps_to_markup(node, [SolveStep("foo.py", insert_node.mark_id, "qux")])
    assert val == "foobarquxbaz"
    val = apply_steps_to_markup(node, [SolveStep("foo.py", "wrong", "qux")])
    assert val == "foobarbaz"
    val = apply_steps_to_markup(
        node,
        [SolveStep("foo.py", insert_node.mark_id, "qux")],
        stop_after_first_change=True,
    )
    assert val == "foobarqux"


def test_fail_module():
    # reference: https://github.com/windyroad/JUnit-Schema/blob/master/JUnit.xsd#L73C7-L73C16
    xml = """<?xml version="1.0" encoding="utf-8"?>
<testsuites>
    <testsuite name="pytest" errors="1" failures="0" skipped="0" tests="1" time="0.044" timestamp="2000-02-27T00:53:12.379941" hostname="530000000a3d">
        <testcase classname="" name="synth_solution" time="0.000">
            <error message="collection failure">ImportError while importing test module
                '/sandbox/synth_solution.py'.
                Hint: make sure your test modules/packages have valid Python names.
                Traceback:
                /usr/local/lib/python3.11/importlib/__init__.py:126: in import_module
                return _bootstrap._gcd_import(name[level:], package, level)
                synth_solution.py:2: in &lt;module&gt;
                from interp.nlptools_custom.nlp_utils import get_wiki_entities, get_wiki_element,
                find_entailment_ratings, get_nlp
                E ModuleNotFoundError: No module named 'interp'
            </error>
        </testcase>
        <testcase classname="" name="synth_solution1" time="0.000">
            <failure message="AssertionError: assert 1 == 0">
    def test_new_sample_no_eval(example_dypy_prob):
        problem_w_doc, function_w_doc = example_dypy_prob
        solver = DummyCodeSolverAutoRegressive(None)
        solution = solver.solve(problem_w_doc, lang_spec=PythonLangSpec())
&gt;       assert len(solution.solve_steps) == 0
E       AssertionError: assert 1 == 0
E        +  where 1 = len((SolveStep(path='project1/grab/document.py', mark_id='tag:StsEditable,4099,4161,1', value='            return unicodec.normalize_encoding_name(encoding)\n'),))
E        +    where (SolveStep(path='project1/grab/document.py', mark_id='tag:StsEditable,4099,4161,1', value='            return unicodec.normalize_encoding_name(encoding)\n'),) = CodeSolution(solve_steps=(SolveStep(path='project1/grab/document.py', mark_id='tag:StsEditable,4099,4161,1', value='  ... instructions_are_essential=False), lang_spec=&lt;synthegrator.lang_spec.PythonLangSpec object at 0x122340f50&gt;, time=None).solve_steps

synthegrator/tests/test_xx_dypybench.py:177: AssertionError
            </failure>
        </testcase>
        <testcase classname="" name="synth_solution2" time="0.000">
            <failure></failure>
        </testcase>
        <testcase classname="" name="synth_solution3" time="0.000">
        </testcase>
    </testsuite>
</testsuites>""".strip()
    results = parse_junit_test_cases(xml, MagicMock())
    assert results.summary.valid_xml
    assert len(results) == 4
    assert not results.all_passed()
    assert not results[0].success
    assert "Hint: " in results[0].fail_text
    assert results[0].fail_message == "collection failure"

    assert not results[1].success
    assert "def test_new_sample_no_eval" in results[1].fail_text
    assert "AssertionError" in results[1].fail_message


def test_invalid_result_xml():
    # reference: https://github.com/windyroad/JUnit-Schema/blob/master/JUnit.xsd#L73C7-L73C16
    xml = """<?xml version="1.0" encoding="utf-8"?>
<testsuites>
    <testsuite name="pytest" errors="1" failures="0" skipped="0" tests="1" time="0.044" timestamp="2000-02-27T00:53:12.379941" hostname="530000000a3d">
        <testcase classname="" name="synth_solution" time="0.000">
            <error message="collection failure">ImportError while importing test module
                '/sandbox/synth_solution.py'.
                Hint: make sure your test modules/packages have valid Python names.
                Traceback:
                /usr/local/lib/python3.11/importlib/__init__.py:126: in import_module
                return _bootstrap._gcd_import(name[level:], package, level)
                synth_solution.py:2: in &lt;module&gt;
                from interp.nlptools_custom.nlp_utils import get_wiki_entities, get_wiki_element,
                find_entailment_ratings, get_nlp
                E ModuleNotFoundError: No module named 'interp'
            </error>
            <failure message="AssertionError: assert 1 == 0">
    def test_new_sample_no_eval(example_dypy_prob):
        problem_w_doc, function_w_doc = example_dypy_prob
        solver = DummyCodeSolverAutoRegressive(None)
        solution = solver.solve(problem_w_doc, lang_spec=PythonLangSpec())
&gt;       assert len(solution.solve_steps) == 0
E       AssertionError: assert 1 == 0
E        +  where 1 = len((SolveStep(path='project1/grab/document.py', mark_id='tag:StsEditable,4099,4161,1', value='            return unicodec.normalize_encoding_name(encoding)\n'),))
E        +    where (SolveStep(path='project1/grab/document.py', mark_id='tag:StsEditable,4099,4161,1', value='            return unicodec.normalize_encoding_name(encoding)\n'),) = CodeSolution(solve_steps=(SolveStep(path='project1/grab/document.py', mark_id='tag:StsEditable,4099,4161,1', value='  ... instructions_are_essential=False), lang_spec=&lt;synthegrator.lang_spec.PythonLangSpec object at 0x122340f50&gt;, time=None).solve_steps

synthegrator/tests/test_xx_dypybench.py:177: AssertionError
            </failure>
        </testcase>
        <testcase classname="" name="synth_solution2" time="0.000">
            <failure></failure>
        </testcase>
        <testcase classname="" name="synth_solution3" time="0.000">
        </testcase>
    </testsuite>
</testsuites>""".strip()
    results = parse_junit_test_cases(xml, MagicMock())
    assert not results.summary.valid_xml
    assert len(results) == 3
    assert not results.all_passed()
    assert not results[0].success
    # This defacto means the last failure/error message will be used!
    assert "def test_new_sample_no_eval" in results[0].fail_text
    assert "AssertionError" in results[0].fail_message

    assert not results[1].success
    assert results[1].fail_text is None
    assert results[1].fail_message is None


def test_discovered_test_cases():
    test_case = DiscoveredTestsuite(
        cmds=[
            # Write some "existing" test suite to a file
            (
                "echo 'import example\ndef test_foo():\n    assert example.foo() == 2'"
                " > /sandbox/test_example.py"
            ),
            # Run pytest on it and get the result
            "python3 -m pytest --junitxml=/sandbox/junit.xml /sandbox/test_example.py",
            "cat /sandbox/junit.xml",
        ],
    )
    problem = CodeProblem(
        working_directory=ProjectWorkingDirectory(
            files=ProjectDir.construct_with_one_file(
                "example.py",
                "def foo():\n    return None",
            ),
        ),
        transformation_spec=TransformationSpec(
            statements=(
                StsSpecStatement(None, StsSelectorLineSpan(1, 2), StsInsert()),
            ),
        ),
        test_cases=[test_case],
    )
    solution = make_solve_for_single_edit(problem, "    return 2")
    result_list = evaluate_code_problem_execution(problem, solution)
    assert len(result_list) == 1
    assert result_list[0].completed
    assert result_list[0].success
    assert result_list.all_passed()


def test_exception_on_bad_solution_indent():
    body = "def foo():\n"
    problem = CodeProblem(
        working_directory=ProjectWorkingDirectory(
            files=ProjectDir.construct_with_one_file("example.py", body),
        ),
        transformation_spec=TransformationSpec(
            statements=(
                StsSpecStatement(
                    None,
                    StsSelectorCharSpan(len(body), len(body)),
                    StsInsert(
                        set_base_space_indent=4,
                    ),
                ),
            ),
        ),
    )
    render = render_single_edit_codeproblem_autoregressive(problem)
    solution = CodeSolution(
        solve_steps=(
            SolveStep("example.py", render.changed_element.mark_id, "    return 2"),
        ),
        problem=problem,
        lang_spec=PythonLangSpec(),
    )
    sol = apply_solution(solution)
    assert sol.get_only_file().content_str == "def foo():\n    return 2"
    solution = CodeSolution(
        solve_steps=(
            SolveStep("example.py", render.changed_element.mark_id, "return 2"),
        ),
        problem=problem,
        lang_spec=PythonLangSpec(),
    )
    with pytest.raises(ValueError):
        apply_solution(solution)


def test_classification_question_eval():
    problem = basic_classification_question_problem()
    classification = get_only_item(
        problem.transformation_spec.classification_question_statements(),
    )
    solution = CodeSolution(
        solve_steps=(
            SolveStep(
                classification.attempt_to_get_path_as_exact_path(),
                classification.static_mark_id,
                "python",
            ),
        ),
        problem=problem,
        lang_spec=PythonLangSpec(),
    )
    problem.known_solutions.append(solution)
    eval = problem.preferred_solution_evaluator.evaluate(solution)
    assert eval.main_metric.is_success
    assert len(eval.extra_metrics) == 3
    assert eval.get_metric_from_key("classification_correct_prob").float_val == 1.0
    solution_2 = CodeSolution(
        solve_steps=(
            SolveStep(
                classification.attempt_to_get_path_as_exact_path(),
                classification.static_mark_id,
                '{"python": 0.1, "java": 0.2, "c++": 0.7}',
            ),
        ),
        problem=problem,
        lang_spec=PythonLangSpec(),
    )
    eval = problem.preferred_solution_evaluator.evaluate(solution_2)
    assert not eval.main_metric.is_success
    assert eval.get_metric_from_key("classification_correct_prob").float_val == 0.1
    print(eval)


def test_exact_match():
    problem = make_simple_method_completion_problem(
        signature="def foo(x) -> int:",
        prompt="Write a function that doubles its input.",
        known_solutions=["    return 4"],
        test_cases=[
            TestCaseMethodCallIsEq(
                test_input=(2,),
                output=4,
                method_name="foo",
                is_hidden_test=False,
                fail_message="",
            ),
            TestCaseMethodCallIsEq(
                test_input=(3,),
                output=6,
                method_name="foo",
                is_hidden_test=False,
                fail_message="",
            ),
        ],
    )
    solve = make_solve_for_single_edit(problem, "    return 4")

    evaluator = problem.preferred_solution_evaluator
    evaluator._extra_metric_funcs = [create_exact_match_metric(strip=True)]

    solve_eval = evaluator.evaluate(solve)

    # one of the actual tests should pass
    assert not solve_eval.main_metric.is_success
    assert solve_eval.main_metric.float_val == 0.5
    assert solve_eval.main_metric.label == TestRunResultList.METRIC_NAME

    # and we expect an exact match
    assert solve_eval.extra_metrics[0].is_success
    assert solve_eval.extra_metrics[0].float_val == 1.0
    assert solve_eval.extra_metrics[0].label == "exact_match"

    # The tests should fail but exact match should succeed
    # due to strip
    solve = make_solve_for_single_edit(problem, "     return 4  ")

    evaluator = problem.preferred_solution_evaluator
    evaluator._extra_metric_funcs = [create_exact_match_metric(strip=True)]

    solve_eval = evaluator.evaluate(solve)

    # the actual tests should pass
    assert not solve_eval.main_metric.is_success
    assert solve_eval.main_metric.float_val == 0.0
    assert solve_eval.main_metric.label == TestRunResultList.METRIC_NAME

    # and we expect an exact match
    assert solve_eval.extra_metrics[0].is_success
    assert solve_eval.extra_metrics[0].float_val == 1.0
    assert solve_eval.extra_metrics[0].label == "exact_match"

    # The tests should pass but exact match should fail
    solve = make_solve_for_single_edit(problem, "    return x*2")

    evaluator = problem.preferred_solution_evaluator
    evaluator._extra_metric_funcs = [create_exact_match_metric(strip=True)]

    solve_eval = evaluator.evaluate(solve)

    # the actual tests should pass
    assert solve_eval.main_metric.is_success
    assert solve_eval.main_metric.float_val == 1.0
    assert solve_eval.main_metric.label == TestRunResultList.METRIC_NAME

    # and we expect an exact match
    assert not solve_eval.extra_metrics[0].is_success
    assert solve_eval.extra_metrics[0].float_val == 0.0
    assert solve_eval.extra_metrics[0].label == "exact_match"

    # Half the tests should pass but exact match should fail
    solve = make_solve_for_single_edit(problem, "    return 4 ")

    evaluator = problem.preferred_solution_evaluator
    evaluator._extra_metric_funcs = [create_exact_match_metric(strip=False)]

    solve_eval = evaluator.evaluate(solve)

    # the actual tests should pass
    assert not solve_eval.main_metric.is_success
    assert solve_eval.main_metric.float_val == 0.5
    assert solve_eval.main_metric.label == TestRunResultList.METRIC_NAME

    # and we expect an exact match
    assert not solve_eval.extra_metrics[0].is_success
    assert solve_eval.extra_metrics[0].float_val == 0.0
    assert solve_eval.extra_metrics[0].label == "exact_match"
