from abc import ABC, abstractmethod
from collections.abc import Sequence

from lmwrapper.structs import LmPrediction

from synthegrator.code_problems import CodeProblem
from synthegrator.lang_specs.lang_spec import LangSpec
from synthegrator.problem_rendering import LmPromptRender, LmPromptRenderSingleEdit
from synthegrator.transformation_spec import (
    MarkElement,
    SolveStep,
    StsEditable,
    StsInsert,
    StsPlaceTransforming,
)
from synthegrator.util import get_first_line_with_nonwhitespace


def format_return_val_for_node(
    value: str,
    node: MarkElement[StsPlaceTransforming],
    lang_spec,
) -> str | tuple[str, tuple[int, int]]:
    """
    Makes sure that a solution value to be put in place of the node
    matches all the node's requirements (like starting with certain indentation
    or whatever).

    If you give it joinable_tokens then will also return a tuple
    representing the best guess of the at the slice of the tokens that
    are actually used in this new value. Note that will not always be an exact
    match. For exampling setting the base indent or setting the ending text
    will not have any analog to the original tokens.
    """
    verb = node.verb
    if not isinstance(verb, StsPlaceTransforming):
        msg = "Verb must be StsPlaceTransforming"
        raise TypeError(msg)
    if not isinstance(verb, StsInsert | StsEditable):
        msg = f"Unhandled verb type: {type(verb)}"
        raise TypeError(msg)

    if verb.set_base_space_indent:
        value = set_base_indent(value.lstrip("\n"), verb.set_base_space_indent)
    if verb.stop_at_block_end:
        value = lang_spec.truncate_function_completion(value)
    if verb.set_ending_text:
        new_start = _get_new_content_start(value, verb.set_ending_text)
        value += verb.set_ending_text[new_start:]
    if verb.max_length_chars is not None:
        raise NotImplementedError
    if verb.max_length_lines is not None:
        lines = value.splitlines(keepends=True)
        value = "".join(lines[: min(verb.max_length_lines, len(lines))])

    return value


class AnswerForSolutionNeverSupported(NotImplementedError):
    """
    Raised when try to get an answer for a solution,
    but it is not implemented for the response parser.
    """


class AnswerForSolutionInvalidForSolve(ValueError):
    """
    Raised when try to get an answer for a solution,
    but there is no known way to construct an answer
    that would produce the solve steps given.
    """


class ResponseParser(ABC):
    def __init__(self):
        pass

    @abstractmethod
    def parse(
        self,
        render: LmPromptRender,
        resp: LmPrediction,
        problem: CodeProblem,
    ) -> list[SolveStep]:
        raise NotImplementedError

    # TODO: consider if can do this more naturally with Protocol/mixin
    def answer_for_solution(
        self,
        render: LmPromptRender,
        problem: CodeProblem,
        solve_steps: Sequence[SolveStep],
    ) -> str:
        raise AnswerForSolutionNeverSupported


class ResponseParserSingleEdit(ResponseParser):
    def __init__(self):
        super().__init__()

    def parse(
        self,
        render: LmPromptRenderSingleEdit,
        resp: LmPrediction,
        problem: CodeProblem,
    ) -> list[SolveStep]:
        lang_spec = problem.get_lang_spec_for_path(render.path)
        if problem.instructions_are_essential or render.prompt.is_text_a_chat():
            solve_steps = pull_solution_from_dialog_completion(
                resp,
                problem,
                lang_spec,
                code_prefix=render.code_prefix,
                insert_node=render.changed_element,
                path=render.path,
            )
        else:
            solution_text = format_return_val_for_node(
                resp.completion_text,
                render.changed_element,
                lang_spec,
            )
            solve_steps = (
                SolveStep(
                    render.path,
                    render.changed_element.mark_id,
                    solution_text,
                ),
            )
        return solve_steps


def set_base_indent(code_text: str, space_indent_level: int) -> str:
    lines = code_text.split("\n")
    first_line = get_first_line_with_nonwhitespace(code_text)
    current_base_indent = len(first_line) - len(first_line.lstrip())
    delta = space_indent_level - current_base_indent

    if delta == 0:
        return code_text

    if delta > 0:
        return "\n".join(
            [" " * delta + line if line.strip() else line for line in lines],
        )

    # delta < 0
    return "\n".join([line[-delta:] for line in lines])


def pull_solution_from_dialog_completion(
    completion: LmPrediction,
    problem: CodeProblem,
    lang_spec: LangSpec,
    code_prefix: str,
    insert_node: MarkElement[StsPlaceTransforming],
    path: str,
) -> list[SolveStep]:
    code = _pull_first_md_block_from_completion(completion)
    if code is None:
        code = completion.completion_text
        if code.endswith("\n```"):
            code = code[: -len("\n```")]
    # Remove docstrings. This is because chat models might choose to add them or not
    code = _strip_prefix_comments(code, lang_spec)
    code_prefix = _strip_prefix_comments(code_prefix, lang_spec)
    code_prefix = code_prefix.rstrip()
    # Try to strip of repeated code from the prompt
    code = code[_get_new_content_start(code_prefix, code) :]
    if code.startswith("\n"):
        code = code[1:]
    code = format_return_val_for_node(code, insert_node, lang_spec)
    if insert_node.verb.stop_at_block_end:
        code = lang_spec.truncate_function_completion(code)
    if lang_spec.get_lang_md_name() == "python" and code.startswith('    """'):
        # Strip docstring
        code = code[code.find('"""') + 3 :]
        code = code[code.find('"""') + 3 :]
        if len(code) > 0 and code[0] == "\n":
            code = code[1:]
    return [SolveStep(path, insert_node.mark_id, code)]


def _strip_prefix_comments(
    code: str,
    lang_spec: LangSpec,
):
    if lang_spec.get_lang_md_name() == "python":
        comments_in_code = list(lang_spec.get_prefix_comments(code))
        for comment in comments_in_code:
            text = "    " + comment.doc_string_node.text.decode() + "\n"
            if text in code:
                code = code.replace(text, "")
    return code


def _pull_first_md_block_from_completion(
    completion: LmPrediction,
) -> str | None:
    """
    Pulls the first markdown block from a completion.
    Markdown blocks are ```...```.
    """
    text = completion.completion_text

    # Find the start of the code block
    start = text.find("```")
    if start == -1:
        return None

    # Find the end of the line after the starting ```
    # This is where the markdown language name is usually located
    end_of_language_name = text.find("\n", start + 3)
    if end_of_language_name == -1:
        # if there's no newline after the starting ```, then it's not a valid block
        return None

    # Now find the end of the code block
    end = text.find("```", end_of_language_name)
    if end == -1:
        # if there's no closing ```, then it's not a valid block
        return None

    # Return the text inside the block, excluding the markdown language name
    return text[end_of_language_name + 1 : end]


def _get_new_content_start(first_str: str, new_content: str) -> int:
    """
    Find the point in the new_content string where to slice so everything
    after is unique from the suffix of the first string.
    This is used to find the point where the new generation starts if
    there is possible repeating of the prompt.

    GPT-4 gen:

    Approach:
        We use the rolling hash technique to compute the hash values of the
        suffix of first_str and the prefix of new_content in linear time.
        While computing these hashes, we compare them at each step.
        If they are equal, it means that the current suffix of first_str
        is also a prefix of new_content. We remember this position.
        The last position at which the hashes were equal will give us the
        longest suffix of first_str that is also a prefix of new_content.

    Note:
    ----
        This function may return false positives (collisions) because two different strings
        might have the same hash, but it is very unlikely when we use a large prime modulus.

    Returns:
    -------
        An integer position value representing the starting point of 'new' content in new_content.

    """
    # TODO: some way of skipping comments

    # Initializing hashes and multipliers
    # m is a large prime number used as modulus to prevent overflow and hash collision
    # p is another prime number used as base of our polynomial hash function
    m, p = 1000000007, 31
    power, hash_first, hash_new = 1, 0, 0

    result = 0  # Initialize result position

    for i in range(len(first_str)):
        # Calculate ASCII value of char
        symbol_first = ord(first_str[-i - 1])
        symbol_new = ord(new_content[i]) if i < len(new_content) else 0

        # Compute the hash of the last (i+1) chars of first_str, and
        # first (i+1) chars of new_content using rolling hash technique.
        hash_first = (hash_first + symbol_first * power) % m
        hash_new = (hash_new * p + symbol_new) % m
        power = (power * p) % m

        # If hashes are equal, update result position
        if hash_first == hash_new:
            result = i + 1

    return result
