from tree_sitter import Node

from synthegrator.util import get_first_line_with_nonwhitespace


def truncate_when_decreases_indent(text: str) -> str:
    first_line = get_first_line_with_nonwhitespace(text)
    starting_indent = len(first_line) - len(first_line.lstrip())
    lines = text.splitlines()
    seen_first_line = False
    for i, line in enumerate(lines):
        if line.lstrip() != "":
            seen_first_line = True
        if (
            len(line) - len(line.lstrip()) < starting_indent
            and seen_first_line
            and line.lstrip() != ""
        ):
            return "\n".join(lines[:i])
    return text


class TreeSitterUtil:
    @staticmethod
    def nodes_are_equal(n1: Node, n2: Node) -> bool:
        if n1 is None or n2 is None:
            return False
        return (
            n1.type == n2.type
            and n1.start_point == n2.start_point
            and n1.end_point == n2.end_point
        )

    @staticmethod
    def node_parent(tree, node) -> Node:
        to_visit = [tree.root_node]
        while len(to_visit) > 0:
            next_node = to_visit.pop()
            for child in next_node.children:
                if TreeSitterUtil.nodes_are_equal(child, node):
                    return next_node
            to_visit.extend(next_node.children)
        msg = "Could not find node in tree."
        raise ValueError(msg)

    @staticmethod
    def copy_tree_excluding_node(node, skip: Node, source_code: str):
        return TreeSitterUtil.copy_tree_excluding_nodes(node, [skip], source_code)

    @staticmethod
    def copy_tree_excluding_nodes(
        node: Node,
        skip: list[Node],
        source_code: str,
    ):
        new_source_code = ""

        # Iterate through the children and recursively call the function.
        child_index = 0
        skipped = False
        for child in node.children:
            if any(TreeSitterUtil.nodes_are_equal(child, _skip) for _skip in skip):
                skipped = True
            else:
                # Add the source code between this child and the previous child (or the start of the parent).
                new_source_code += source_code[
                    node.start_byte + child_index : child.start_byte
                ]

                # Recursively handle the child.
                new_source_code += TreeSitterUtil.copy_tree_excluding_nodes(
                    child,
                    skip,
                    source_code,
                )

            # Update the index to the end of this child.
            child_index = child.end_byte - node.start_byte

        # Add the source code after the last child.
        if skipped:
            new_source_code = new_source_code.lstrip()

        new_source_code += source_code[node.start_byte + child_index : node.end_byte]

        return new_source_code
