from unittest.mock import MagicMock

from synthegrator.problem_rendering_insertion_tags import (
    TaggedEditRenderer,
    TaggedEditResponseParser,
    _find_ids_and_values_in_tagged_answer,
    _get_mark_default_id,
    _marked_up_text_to_tagged_prompt_text,
)
from synthegrator.tests.non_dataset_tests.test_problem_rendering import (
    factor_problems,
)
from synthegrator.transformation_spec import (
    MarkElement,
    MarkText,
    StsInsert,
    StsInvisible,
    StsNormalText,
)

assert factor_problems, "Unused import detectors struggle with fixtures"


def test_tagged_prompt():
    markup = MarkElement(
        verb=StsNormalText(),
        children=[
            MarkText("foo", 0, 3),
            MarkElement(StsInsert(), [], 3, 3),
            MarkText("bar", 3, 6),
        ],
        start_char=0,
        end_char=6,
    ).freeze()
    out, tag_ids = _marked_up_text_to_tagged_prompt_text(
        markup,
        path_id="0",
        # path="solution.py",
    )
    edit_node: MarkElement = markup.children[1]
    assert out == f'foo<edit id="0_{edit_node.mark_id}"></edit>bar'
    assert tag_ids == [f"0_{edit_node.mark_id}"]


def test_tagged_prompt2():
    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=6,
    ).freeze()
    out, _tag_ids = _marked_up_text_to_tagged_prompt_text(
        markup,
        path_id="0",
    )
    markup.children[1]
    assert out == "foobar"


def test_insert_prompt(factor_problems):
    problem = factor_problems[0]
    renderer = TaggedEditRenderer()
    render = renderer(problem, MagicMock())
    content = (
        "Complete the following method please\n"
        "@@ foo.py @@\n"
        "def double(x) -> int:\n"
        '    """\n'
        "    Write a function that doubles its input.\n"
        '    """\n'
        '<edit id="0_qkIO0"></edit>\n'
        "\n"
        "---\n"
        "Implement the correct changes in the above code in the <edit> tags. Please "
        'place each fix in a separate <edit_solve id="id"> tags.\n'
        "For example:\n"
        '<edit_solve id="id_example">new replacement for the text...</edit_solve>\n'
        "Answer:\n"
    )
    assert render.prompt.text == content


def test_insert_prompt_default_id(factor_problems):
    problem = factor_problems[0]
    default = _get_mark_default_id(problem)
    assert default == "0_qkIO0"


def test_insert_prompt_custom(factor_problems):
    problem = factor_problems[0]
    renderer = TaggedEditRenderer(
        tag_name_edit="buggy",
        tag_name_solve="fix",
        custom_closing_lines="Please fix the code.",
    )
    render = renderer(problem, MagicMock())
    content = (
        "Complete the following method please\n"
        "@@ foo.py @@\n"
        "def double(x) -> int:\n"
        '    """\n'
        "    Write a function that doubles its input.\n"
        '    """\n'
        '<buggy id="0_qkIO0"></buggy>\n'
        "\n"
        "---\n"
        "Please fix the code."
    )
    assert render.prompt.text == content
    response_parser = TaggedEditResponseParser()
    parse = response_parser.parse(
        render,
        MagicMock(
            completion_text='<fix id="0_qkIO0">    return x * 2</fix>',
        ),
        problem,
    )
    assert len(parse) == 1
    assert parse[0].path == "foo.py"
    assert parse[0].mark_id == "qkIO0"
    assert parse[0].value == "    return x * 2"


def test_find_ids_in_answer(factor_problems):
    val = _find_ids_and_values_in_tagged_answer(
        'foo<edit_solve id="0_qkIO0">bar</edit_solve>baz',
        default_id="default",
    )
    assert val == {"0_qkIO0": "bar"}


def test_find_ids_in_answer_dup(factor_problems):
    val = _find_ids_and_values_in_tagged_answer(
        '<edit id="0_qkIO0">foo</edit><edit_solve id="0_qkIO0">bar</edit>',
        default_id="default",
    )
    assert val == {"0_qkIO0": "bar"}


def test_find_ids_in_answer_many(factor_problems):
    val = _find_ids_and_values_in_tagged_answer(
        '<edit_solve id="a">foo</edit_solve><edit_solve id="0_qkIO0">bar</edit>',
        default_id="default",
    )
    assert val == {"a": "foo", "0_qkIO0": "bar"}


def test_find_ids_in_answer_stop(factor_problems):
    val = _find_ids_and_values_in_tagged_answer(
        'foo\n<edit_solve id="0_qkIO0">bar',
        default_id="default",
    )
    assert val == {"0_qkIO0": "bar"}


def test_find_ids_in_answer_default(factor_problems):
    val = _find_ids_and_values_in_tagged_answer("foo", default_id="default")
    assert val == {"default": "foo"}


def test_find_ids_in_answer_default(factor_problems):
    val = _find_ids_and_values_in_tagged_answer(
        "prestuff<edit_solve>foo</edit_solve>", default_id="defaultid"
    )
    assert val == {"defaultid": "foo"}


def test_find_ids_in_answer_default2(factor_problems):
    val = _find_ids_and_values_in_tagged_answer(
        "foo</edit_solve>bar", default_id="default"
    )
    assert val == {"default": "foo"}


def test_full_parser(factor_problems):
    parser = TaggedEditResponseParser()
    problem = factor_problems[0]
    renderer = TaggedEditRenderer()
    render = renderer(problem, MagicMock())
    solve_steps = parser.parse(
        render,
        MagicMock(
            completion_text='<edit_solve id="0_qkIO0">    return x * 2</edit_solve>',
        ),
        problem,
    )
    assert len(solve_steps) == 1
    assert solve_steps[0].path == "foo.py"
    assert solve_steps[0].mark_id == "qkIO0"
    assert solve_steps[0].value == "    return x * 2"
