from unittest.mock import MagicMock

import pytest
from lmwrapper.structs import ChatGptRoles

from synthegrator.code_problem_builders import make_simple_method_completion_problem
from synthegrator.code_problems import (
    TestCaseMethodCallIsEq,
)
from synthegrator.few_shotting import FewShotConfig, FewShotLibrary
from synthegrator.problem_rendering import (
    render_markup_to_text_autoregressive,
    render_single_edit_codeproblem_dialog,
)
from synthegrator.prompting_test_case_selection import (
    RandomVisibleTestCaseSelectionStrategy,
)
from synthegrator.tests.non_dataset_tests.test_java_stuff import get_java_double_problem
from synthegrator.transformation_spec import (
    MarkElement,
    MarkText,
    StsInsert,
    StsInvisible,
    StsNormalText,
)


def test_markup_autoregressive_rendering():
    markup = MarkElement(
        verb=StsNormalText(),
        children=[
            MarkText("foo", 0, 3),
        ],
        start_char=0,
        end_char=3,
    )
    problem = MagicMock()
    problem.instructions_are_essential = False
    out = render_markup_to_text_autoregressive(
        markup,
        path="solution.py",
        problem=problem,
    )
    assert out.prompt.text == "foo"
    assert out.path == "solution.py"
    assert out.changed_element is None


def test_markup_autoregressive_rendering2():
    markup = MarkElement(
        verb=StsNormalText(),
        children=[
            MarkText("foo", 0, 3),
            MarkElement(StsInsert(), [], 3, 3),
            MarkText("bar", 3, 6),
        ],
        start_char=0,
        end_char=6,
    )
    problem = MagicMock()
    problem.instructions_are_essential = False
    out = render_markup_to_text_autoregressive(
        markup,
        path="solution.py",
        problem=problem,
    )
    assert out.prompt.text == "foo"
    assert out.path == "solution.py"
    assert out.changed_element == markup.children[1]


def test_markup_autoregressive_rendering3():
    markup = MarkElement(
        verb=StsNormalText(),
        children=[
            MarkText("foo", 0, 3),
            MarkElement(StsInvisible(), [MarkText("hide", 3, 7)], 3, 7),
            MarkText("bar", 7, 10),
        ],
        start_char=0,
        end_char=10,
    )
    problem = MagicMock()
    problem.instructions_are_essential = False
    out = render_markup_to_text_autoregressive(
        markup,
        path="solution.py",
        problem=problem,
    )
    assert out.prompt.text == "foobar"
    assert out.path == "solution.py"
    assert out.changed_element is None


def make_test_case(test_input, output, method_name):
    return TestCaseMethodCallIsEq(
        test_input=test_input,
        output=output,
        method_name=method_name,
        fail_message="",
        is_hidden_test=False,
    )


@pytest.fixture
def factor_problems():
    return [
        make_simple_method_completion_problem(
            signature="def double(x) -> int:",
            prompt="Write a function that doubles its input.",
            known_solutions=[
                "    return x * 2\n",
            ],
            test_cases=[
                make_test_case((2,), 4, "double"),
                make_test_case((3,), 6, "double"),
            ],
            target_path="foo.py",
            override_instructions="Complete the following method please",
            problem_id="double",
        ),
        make_simple_method_completion_problem(
            signature="def triple(x) -> int:",
            prompt="Write a function that triples its input.",
            known_solutions=[
                "    return x * 3\n",
            ],
            test_cases=[
                make_test_case((2,), 6, "triple"),
                make_test_case((3,), 9, "triple"),
            ],
            target_path="foo.py",
            override_instructions="Complete the following method please",
            problem_id="triple",
        ),
        make_simple_method_completion_problem(
            signature="def quadruple(x) -> int:",
            prompt="Write a function that quadruples its input.",
            known_solutions=[
                "    return x * 4\n",
            ],
            test_cases=[
                make_test_case((2,), 8, "quadruple"),
                make_test_case((3,), 12, "quadruple"),
            ],
            target_path="foo.py",
            override_instructions="Complete the following method please",
            problem_id="quadruple",
        ),
    ]


@pytest.fixture
def factor_problems_java():
    return [get_java_double_problem()]


def test_dialog_rendering(factor_problems):
    problem = factor_problems[0]
    render = render_single_edit_codeproblem_dialog(problem)
    assert render.path == "foo.py"
    assert len(render.prompt.text) == 1
    content = (
        "Complete the following method please\n"
        "```\n"
        "def double(x) -> int:\n"
        '    """\n'
        "    Write a function that doubles its input.\n"
        '    """\n'
        "```\n"
    )
    assert render.prompt.text[0].content == content
    assert render.prompt.text[0].role == ChatGptRoles.user
    element = render.changed_element
    assert isinstance(element, MarkElement)
    assert element.verb == StsInsert(
        set_base_space_indent=4,
        stop_at_block_end=True,
    )
    assert element.mark_end_original_char == len(
        "def double(x) -> int:\n"
        '    """\n'
        "    Write a function that doubles its input.\n"
        '    """\n',
    )


def test_dialog_rendering_few_shot(factor_problems):
    problem = factor_problems[0]
    render = render_single_edit_codeproblem_dialog(
        problem,
        few_shot_config=FewShotConfig(
            library=FewShotLibrary(factor_problems[:2]),
            num_examples=1,
            change_examples_between_solves=False,
        ),
        prompt_seed=0,
    )
    assert render.path == "foo.py"
    assert len(render.prompt.text) == 3
    content = (
        "Complete the following method please\n"
        "```\n"
        "def double(x) -> int:\n"
        '    """\n'
        "    Write a function that doubles its input.\n"
        '    """\n'
        "```\n"
    )
    assert render.prompt.text[2].content == content
    element = render.changed_element
    assert isinstance(element, MarkElement)
    assert element.verb == StsInsert(
        set_base_space_indent=4,
        stop_at_block_end=True,
    )
    assert element.mark_end_original_char == len(
        "def double(x) -> int:\n"
        '    """\n'
        "    Write a function that doubles its input.\n"
        '    """\n',
    )
    assert (
        render.prompt.text[0].content == "Complete the following method please\n"
        "```\n"
        "def triple(x) -> int:\n"
        '    """\n'
        "    Write a function that triples its input.\n"
        '    """\n'
        "```\n"
    )
    assert render.prompt.text[1].content == "```\n    return x * 3\n```"
    assert render.prompt.text[0].role == ChatGptRoles.user
    assert render.prompt.text[1].role == ChatGptRoles.assistant
    assert render.prompt.text[2].role == ChatGptRoles.user


def test_dialog_rendering_few_shot_test_case_prompt(factor_problems):
    problem = factor_problems[0]
    render = render_single_edit_codeproblem_dialog(
        problem,
        few_shot_config=FewShotConfig(
            library=FewShotLibrary(factor_problems[:2]),
            num_examples=1,
            change_examples_between_solves=False,
        ),
        prompt_seed=0,
        prompt_test_case_selection_strategy=RandomVisibleTestCaseSelectionStrategy(0),
    )
    assert render.path == "foo.py"
    assert len(render.prompt.text) == 3
    content = (
        "Complete the following method please\n"
        "```\n"
        "def double(x) -> int:\n"
        '    """\n'
        "    Write a function that doubles its input.\n"
        '    """\n'
        "```\n"
        "Test cases:\n"
        "```\n"
        "def test_double_0():\n"
        "    assert double(2) == 4\n"
        "def test_double_1():\n"
        "    assert double(3) == 6\n"
        "```"
    )
    assert render.prompt.text[2].content == content
    element = render.changed_element
    assert isinstance(element, MarkElement)
    assert element.verb == StsInsert(
        set_base_space_indent=4,
        stop_at_block_end=True,
    )
    assert element.mark_end_original_char == len(
        "def double(x) -> int:\n"
        '    """\n'
        "    Write a function that doubles its input.\n"
        '    """\n',
    )
    assert (
        render.prompt.text[0].content == "Complete the following method please\n"
        "```\n"
        "def triple(x) -> int:\n"
        '    """\n'
        "    Write a function that triples its input.\n"
        '    """\n'
        "```\n"
        "Test cases:\n"
        "```\n"
        "def test_triple_0():\n"
        "    assert triple(2) == 6\n"
        "def test_triple_1():\n"
        "    assert triple(3) == 9\n"
        "```"
    )
    assert render.prompt.text[1].content == "```\n    return x * 3\n```"
    assert render.prompt.text[0].role == ChatGptRoles.user
    assert render.prompt.text[1].role == ChatGptRoles.assistant
    assert render.prompt.text[2].role == ChatGptRoles.user
