# Copyright 2017-2020 Palantir Technologies, Inc.
# Copyright 2021- Python Language Server Contributors.

import re

import parso
import parso.python.tree as tree_nodes

from pylsp import hookimpl

SKIP_NODES = (tree_nodes.Module, tree_nodes.IfStmt, tree_nodes.TryStmt)
IDENTATION_REGEX = re.compile(r'(\s+).+')


@hookimpl
def pylsp_folding_range(document):
    program = document.source + '\n'
    lines = program.splitlines()
    tree = parso.parse(program)
    ranges = __compute_folding_ranges(tree, lines)

    results = []
    for (start_line, end_line) in ranges:
        start_line -= 1
        end_line -= 1
        # If start/end character is not defined, then it defaults to the
        # corresponding line last character
        results.append({
            'startLine': start_line,
            'endLine': end_line,
        })
    return results


def __merge_folding_ranges(left, right):
    for start in list(left.keys()):
        right_start = right.pop(start, None)
        if right_start is not None:
            left[start] = max(right_start, start)
    left.update(right)
    return left


def __empty_identation_stack(identation_stack, level_limits,
                             current_line, folding_ranges):
    while identation_stack != []:
        upper_level = identation_stack.pop(0)
        level_start = level_limits.pop(upper_level)
        folding_ranges.append((level_start, current_line))
    return folding_ranges


def __match_identation_stack(identation_stack, level, level_limits,
                             folding_ranges, current_line):
    upper_level = identation_stack.pop(0)
    while upper_level >= level:
        level_start = level_limits.pop(upper_level)
        folding_ranges.append((level_start, current_line))
        upper_level = identation_stack.pop(0)
    identation_stack.insert(0, upper_level)
    return identation_stack, folding_ranges


def __compute_folding_ranges_identation(text):
    lines = text.splitlines()
    folding_ranges = []
    identation_stack = []
    level_limits = {}
    current_level = 0
    current_line = 0
    while lines[current_line] == '':
        current_line += 1
    for i, line in enumerate(lines):
        if i < current_line:
            continue
        i += 1
        identation_match = IDENTATION_REGEX.match(line)
        if identation_match is not None:
            whitespace = identation_match.group(1)
            level = len(whitespace)
            if level > current_level:
                level_limits[current_level] = current_line
                identation_stack.insert(0, current_level)
                current_level = level
            elif level < current_level:
                identation_stack, folding_ranges = __match_identation_stack(
                    identation_stack, level, level_limits, folding_ranges,
                    current_line)
                current_level = level
        else:
            folding_ranges = __empty_identation_stack(
                identation_stack, level_limits, current_line, folding_ranges)
            current_level = 0
        if line.strip() != '':
            current_line = i
    folding_ranges = __empty_identation_stack(
        identation_stack, level_limits, current_line, folding_ranges)
    return dict(folding_ranges)


def __check_if_node_is_valid(node):
    valid = True
    if isinstance(node, tree_nodes.PythonNode):
        kind = node.type
        valid = kind not in {'decorated', 'parameters', 'dictorsetmaker',
                             'testlist_comp'}
        if kind == 'suite':
            if isinstance(node.parent, tree_nodes.Function):
                valid = False
    return valid


def __handle_skip(stack, skip):
    body = stack[skip]
    children = [body]
    if hasattr(body, 'children'):
        children = body.children
    stack = stack[:skip] + children + stack[skip + 1:]
    node = body
    end_line, _ = body.end_pos
    return node, end_line


def __handle_flow_nodes(node, end_line, stack):
    from_keyword = False
    if isinstance(node, tree_nodes.Keyword):
        from_keyword = True
        if node.value in {'if', 'elif', 'with', 'while'}:
            node, end_line = __handle_skip(stack, 2)
        elif node.value in {'except'}:
            first_node = stack[0]
            if isinstance(first_node, tree_nodes.Operator):
                node, end_line = __handle_skip(stack, 1)
            else:
                node, end_line = __handle_skip(stack, 2)
        elif node.value in {'for'}:
            node, end_line = __handle_skip(stack, 4)
        elif node.value in {'else'}:
            node, end_line = __handle_skip(stack, 1)
    return end_line, from_keyword, node, stack


def __compute_start_end_lines(node, stack):
    start_line, _ = node.start_pos
    end_line, _ = node.end_pos
    modified = False
    end_line, from_keyword, node, stack = __handle_flow_nodes(
        node, end_line, stack)

    last_leaf = node.get_last_leaf()
    last_newline = isinstance(last_leaf, tree_nodes.Newline)
    last_operator = isinstance(last_leaf, tree_nodes.Operator)
    node_is_operator = isinstance(node, tree_nodes.Operator)
    last_operator = last_operator or not node_is_operator

    end_line -= 1

    if isinstance(node.parent, tree_nodes.PythonNode) and not from_keyword:
        kind = node.type
        if kind in {'suite', 'atom', 'atom_expr', 'arglist'}:
            if len(stack) > 0:
                next_node = stack[0]
                next_line, _ = next_node.start_pos
                if next_line > end_line:
                    end_line += 1
                    modified = True
    if not last_newline and not modified and not last_operator:
        end_line += 1
    return start_line, end_line, stack


def __compute_folding_ranges(tree, lines):
    folding_ranges = {}
    stack = [tree]

    while len(stack) > 0:
        node = stack.pop(0)
        if isinstance(node, tree_nodes.Newline):
            # Skip newline nodes
            continue
        if isinstance(node, tree_nodes.PythonErrorNode):
            # Fallback to indentation-based (best-effort) folding
            start_line, _ = node.start_pos
            start_line -= 1
            padding = [''] * start_line
            text = '\n'.join(padding + lines[start_line:]) + '\n'
            identation_ranges = __compute_folding_ranges_identation(text)
            folding_ranges = __merge_folding_ranges(
                folding_ranges, identation_ranges)
            break
        if not isinstance(node, SKIP_NODES):
            valid = __check_if_node_is_valid(node)
            if valid:
                start_line, end_line, stack = __compute_start_end_lines(
                    node, stack)
                if end_line > start_line:
                    current_end = folding_ranges.get(start_line, -1)
                    folding_ranges[start_line] = max(current_end, end_line)
        if hasattr(node, 'children'):
            stack = node.children + stack

    folding_ranges = sorted(folding_ranges.items())
    return folding_ranges