"""
Some functions that can help make common types of problems.
TODO: a lot of these are pretty unweildy. There should
    probably be something with a builder pattern to
    gradually add different parts of the problem.
"""

import dataclasses

from synthegrator.code_problems import (
    CodeProblem,
    CodeSolution,
    LmCodeSolution,
    TestCase,
)
from synthegrator.environments import CodeExecutionEnvironment, ProjectWorkingDirectory
from synthegrator.lang_specs.lang_spec import LangSpec
from synthegrator.lang_specs.lang_spec_python import (
    PythonLangSpec,
    assert_spec_is_python,
)
from synthegrator.memory_fs import ProjectDir
from synthegrator.sandboxing import DockerExecutionContext
from synthegrator.transformation_spec import (
    MarkElement,
    SolveStep,
    StsEditable,
    StsInsert,
    StsInvisible,
    StsPath,
    StsSelectorCharSpan,
    StsSelectorLineSpan,
    StsSpecStatement,
    TransformationSpec,
    markup_path,
)


def make_simple_method_completion_problem(
    prompt: str,
    signature: str,
    test_cases: list[TestCase],
    model_context: str = "",
    execution_context_hidden_from_model: str = "",
    target_path: str = "solution.py",
    known_solutions: list[str] | None = None,
    lang_spec: LangSpec | None = None,
    prompt_already_docstring: bool = False,
    dataset_name: str | None = None,
    problem_id: str | None = None,
    docker_env: DockerExecutionContext = None,
    override_instructions: str | None = None,
    base_indent: int = 4,
    after_method_src: str | None = None,
) -> CodeProblem:
    if lang_spec is None:
        lang_spec = PythonLangSpec()
    assert_spec_is_python(lang_spec)
    if "\r" in signature:
        msg = "Signature not expected to have carriage returns"
        raise ValueError(msg)
    if "\r" in prompt:
        msg = "Prompt not expected to have carriage returns"
        raise ValueError(msg)
    signature = signature.rstrip("\n")
    signature_last_line = signature.split("\n")[-1]
    line_delta = len(signature_last_line) - len(signature_last_line.lstrip())
    indentation = " " * (line_delta + 4)
    if not prompt_already_docstring:
        prompt = (
            indentation
            + '"""\n'
            + "\n".join(f"{indentation}{line}" for line in prompt.split("\n"))
            + f"\n{indentation}"
            + '"""'
        )
    body = (
        execution_context_hidden_from_model
        + model_context
        + signature
        + ("\n" if not signature.endswith("\n") else "")
        + prompt
        + ("\n" if not prompt.endswith("\n") else "")
    )
    post_context = None if (not after_method_src) else "\n" + after_method_src
    return _auto_regressive_insert_problem(
        body,
        override_instructions
        or (
            "Complete the following method. "
            "Give only the solution body, and do not repeat the signature or docstring."
        ),
        target_path,
        test_cases,
        (
            [(0, len(execution_context_hidden_from_model))]
            if execution_context_hidden_from_model
            else None
        ),
        dataset_name,
        problem_id,
        known_solutions,
        lang_spec,
        docker_env,
        base_indent,
        stop_at_block_end=True,
        post_context=post_context,
    )


def _solution_strings_to_solves(
    markup: MarkElement,
    solution_strings: list[str],
    target_path: str,
    problem: CodeProblem,
    lang_spec: LangSpec,
) -> list[CodeSolution]:
    solution_solves = []
    for sol in solution_strings:
        insert_nodes = [
            node
            for node in markup.depth_first_iterate()
            if isinstance(node, MarkElement) and isinstance(node.verb, StsInsert)
        ]
        if len(insert_nodes) != 1:
            msg = "Insert node length must be one"
            raise ValueError(msg)
        insert_node = insert_nodes[0]
        solution_solves.append(
            CodeSolution(
                solve_steps=[
                    SolveStep(
                        path=target_path,
                        mark_id=insert_node.mark_id,
                        value=sol,
                    ),
                ],
                problem=problem,
                lang_spec=lang_spec,
            ),
        )
    return solution_solves


def _auto_regressive_insert_problem(
    body: str,
    instructions: str,
    target_path: str,
    test_cases: list[TestCase],
    invisible_spans: list[tuple[int, int]] | None = None,
    dataset_name: str | None = None,
    problem_id: str | None = None,
    known_solutions: list[str] | None = None,
    lang_spec: LangSpec | None = None,
    docker_env: DockerExecutionContext = None,
    base_indent: int = 4,
    stop_at_block_end: bool = True,
    post_context: str | None = None,
) -> CodeProblem:
    """Makes a problem with a single insert at the end"""
    if lang_spec is None:
        lang_spec = PythonLangSpec()
    if known_solutions is not None and type(known_solutions) != list:
        msg = "known_solutions must be a list"
        raise TypeError(msg)

    if post_context is None:
        post_context = ""

    wd = ProjectWorkingDirectory(
        files=ProjectDir.construct_with_one_file(target_path, body + post_context),
    )
    steps = [
        StsSpecStatement(
            StsPath(target_path),
            StsSelectorCharSpan(len(body), len(body)),
            StsInsert(
                set_base_space_indent=base_indent,
                stop_at_block_end=stop_at_block_end,
            ),
        ),
    ]
    steps.extend(
        StsSpecStatement(None, StsSelectorCharSpan(start, end), StsInvisible())
        for start, end in (invisible_spans or [])
    )
    spec = TransformationSpec(tuple(steps))
    markup = markup_path(wd, target_path, spec)
    problem = CodeProblem(
        instructions=instructions,
        working_directory=wd,
        transformation_spec=spec,
        test_cases=test_cases,
        dataset_name=dataset_name,
        problem_id=problem_id,
        environment=CodeExecutionEnvironment(docker_env=docker_env),
        instructions_are_essential=False,
    )
    problem.known_solutions = _solution_strings_to_solves(
        markup,
        known_solutions or [],
        target_path,
        problem,
        lang_spec,
    )
    return problem


def make_simple_file_generation(
    prompt: str,
    test_cases: list[TestCase],
    model_context: str = "",
    execution_context_hidden_from_model: str = "",
    target_path: str = "solution.py",
    lang_spec: LangSpec | None = None,
    dataset_name: str | None = None,
    problem_id: str | None = None,
    known_solutions: list[str] | None = None,
    docker_env: DockerExecutionContext = None,
) -> CodeProblem:
    if lang_spec is None:
        lang_spec = PythonLangSpec()
    # TODO: We are currently just quickly porting from old method, but
    #    we should support specifying file generation with a spec statement
    assert_spec_is_python(lang_spec)
    prompt = '""""\n' + "    \n".join(line for line in prompt.split("\n")) + '\n"""'
    content = (
        execution_context_hidden_from_model
        + "\n"
        + model_context
        + "\n"
        + prompt
        + "\n"
    )
    return _auto_regressive_insert_problem(
        body=content,
        instructions=f"Write {lang_spec.get_lang_md_name()} code following the instructions below",
        target_path=target_path,  # Fixed: this was previously swapped with instructions
        test_cases=test_cases,
        invisible_spans=(
            [(0, len(execution_context_hidden_from_model))]
            if execution_context_hidden_from_model
            else None
        ),
        dataset_name=dataset_name,
        problem_id=problem_id,
        known_solutions=known_solutions,
        lang_spec=lang_spec,
        docker_env=docker_env,
        base_indent=0,
    )


def _estimate_base_indent(text: str) -> int:
    """Estimates the base indentation of a text"""
    left_ws_len = len(text) - len(text.lstrip())
    left_ws = text[:left_ws_len]
    left_ws.replace("\t", " " * 4)
    left_ws.replace("\n", "")
    left_ws.replace("\r", "")
    return len(left_ws)


def make_simple_line_edit_problem(
    code: str,
    edit_lines_span: tuple[int, int],
    test_cases: list[TestCase],
    invisible_spans: list[tuple[int, int]] | None = None,
    current_text_visible: bool = True,
    target_path: str = "solution.py",
    known_solutions: list[str] | None = None,
    lang_spec: LangSpec | None = None,
    dataset_name: str | None = None,
    problem_id: str | None = None,
    docker_env: DockerExecutionContext = None,
    override_instructions: str | None = None,
    cap_to_max_same_num_lines: bool = True,
) -> CodeProblem:
    if lang_spec is None:
        lang_spec = PythonLangSpec()
    # assert_spec_is_python(lang_spec)
    start_idx, end_idx = edit_lines_span
    base_indent = _estimate_base_indent(code.splitlines()[start_idx])
    edit_statement = StsSpecStatement(
        StsPath(target_path),
        StsSelectorLineSpan(
            start_idx,
            end_idx,
            include_trailing_newline=True,
        ),
        StsEditable(
            current_text_visible=current_text_visible,
            set_ending_text="\n",
            set_base_space_indent=base_indent,
            max_length_lines=(
                end_idx - start_idx if cap_to_max_same_num_lines else None
            ),
        ),
    )

    if invisible_spans is None:
        statements = (edit_statement,)
    else:
        invisible_statements = [
            StsSpecStatement(
                StsPath(target_path),
                StsSelectorLineSpan(invisible_span_start, invisible_span_end),
                StsInvisible(),
            )
            for invisible_span_start, invisible_span_end in invisible_spans
        ]
        statements = (*invisible_statements, edit_statement)

    problem = CodeProblem(
        instructions=override_instructions,  # or "Edit the marked lines",
        working_directory=ProjectWorkingDirectory(
            files=ProjectDir.construct_with_one_file(target_path, code),
        ),
        transformation_spec=TransformationSpec(statements),
        test_cases=test_cases,
        dataset_name=dataset_name,
        problem_id=problem_id,
        environment=CodeExecutionEnvironment(docker_env=docker_env),
        instructions_are_essential=False,
    )
    if (
        problem.transformation_spec.count_editing_statements() != 1
    ):  # TODO only one StsPlaceTransforming
        msg = "Only one editing statement allowed"
        raise ValueError(msg)

    problem.known_solutions = [
        make_solve_for_single_edit(problem, sol) for sol in (known_solutions or [])
    ]
    return problem


def make_solve_for_single_edit(problem, solution: str) -> CodeSolution:
    from synthegrator.problem_rendering import (
        render_single_edit_codeproblem_autoregressive,
    )

    render = render_single_edit_codeproblem_autoregressive(problem)
    return LmCodeSolution(
        solve_steps=(SolveStep(render.path, render.changed_element.mark_id, solution),),
        problem=problem,
        lang_spec=PythonLangSpec(),
        prompt=dataclasses.replace(render.prompt, max_tokens=500),
    )


def make_instruction_only_problem(
    instructions: str,
    insert_node: StsInsert,
    test_cases: list[TestCase],
    known_solutions: list[str],
    dataset_name: str | None = None,
    problem_id: str | None = None,
    target_file_name: str = "solution.py",
) -> CodeProblem:
    wd, transformation_spec, insert_node_element = _instruction_only_data(
        target_file_name,
        insert_node,
    )
    problem = CodeProblem(
        instructions=instructions,
        working_directory=wd,
        test_cases=test_cases,
        transformation_spec=transformation_spec,
        dataset_name=dataset_name,
        problem_id=problem_id,
    )
    problem.known_solutions = [
        CodeSolution(
            solve_steps=(
                SolveStep(
                    target_file_name,
                    insert_node_element.mark_id,
                    known_sol,
                ),
            ),
            problem=problem,
            lang_spec=PythonLangSpec(),
        )
        for known_sol in known_solutions
    ]
    return problem


_wd_cache = {}


def _instruction_only_data(
    target_file_name: str,
    insert_node: StsInsert,
) -> tuple[ProjectWorkingDirectory | None, TransformationSpec, MarkElement]:
    wd = _wd_cache.get(target_file_name)
    if wd is None:
        wd = ProjectWorkingDirectory(
            ProjectDir.construct_with_one_file(target_file_name, ""),
        )
        _wd_cache[target_file_name] = wd
    transformation_spec = TransformationSpec(
        (
            StsSpecStatement(
                path_selector=StsPath(target_file_name),
                text_selector=StsSelectorCharSpan(0, 0),
                verb=insert_node,
            ),
        ),
    )
    markup = markup_path(wd, target_file_name, spec=transformation_spec)
    insert_node_element = next(
        node
        for node in markup.depth_first_iterate()
        if isinstance(node, MarkElement) and isinstance(node.verb, StsInsert)
    )
    return wd, transformation_spec, insert_node_element


"""
TODO?
class ProblemBuilder:
    def __init__(self):
        self._files = {}

    def add_file(self, file_name: str, content: str):
        pass

    def build(self):
        return CodeProblem(

        )
"""
