# Copyright (c) Facebook, Inc. and its affiliates.
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.

from dataclasses import dataclass
from typing import Union

import libcst as cst
from libcst.metadata import (
    CodePosition,
    ExperimentalReentrantCodegenProvider,
    MetadataWrapper,
    ParentNodeProvider,
    WhitespaceInclusivePositionProvider,
)


@dataclass(frozen=True)
class LintPatch:
    """
    Represents the diff generated by an auto-fixer for arc lint.
    """

    # The point in source where the replacement was made. Typically this is the position
    # of the replaced node. Available as both a byte offset and a position.
    start_offset: int
    start_position: CodePosition
    # diff text
    original_diff_str: str
    patched_diff_str: str

    @staticmethod
    def get(
        wrapper: MetadataWrapper,
        original_node: cst.CSTNode,
        replacement_node: Union[cst.CSTNode, cst.FlattenSentinel, cst.RemovalSentinel],
    ) -> "LintPatch":
        # Batch the execution of these position providers
        wrapper.resolve_many(
            [
                ParentNodeProvider,
                ExperimentalReentrantCodegenProvider,
                WhitespaceInclusivePositionProvider,
            ]
        )

        # Use the resolve() API to fetch the data, because it's typed better than
        # resolve_many() is.
        parents = wrapper.resolve(ParentNodeProvider)
        positions = wrapper.resolve(WhitespaceInclusivePositionProvider)
        codegen_partials = wrapper.resolve(ExperimentalReentrantCodegenProvider)

        if isinstance(original_node, cst.Module) and isinstance(
            replacement_node, cst.RemovalSentinel
        ):
            raise Exception("Removing the entire module is not possible")

        # The reentrant codegen provider can only rewrite entire statements at a time,
        # so we need to inspect our parents until find a statement or the module
        possible_statement = original_node
        if isinstance(replacement_node, cst.RemovalSentinel):
            # reentrant codegen doesn't support RemovalSentinel, so use the parent instead
            possible_statement = parents[possible_statement]
        while True:
            if possible_statement in codegen_partials:
                partial = codegen_partials[possible_statement]
                patched_statement = cst.ensure_type(
                    _replace_or_remove(
                        possible_statement, original_node, replacement_node
                    ),
                    cst.BaseStatement,
                )
                original_str = partial.get_original_statement_code()
                patched_str = partial.get_modified_statement_code(patched_statement)
                return LintPatch(
                    partial.start_offset,
                    positions[possible_statement].start,
                    original_str,
                    patched_str,
                )
            elif possible_statement in parents:
                possible_statement = parents[possible_statement]
            else:
                # There's no more parents, so we have to fall back to replacing the whole
                # module.
                original_str = wrapper.module.code
                patched_module = cst.ensure_type(
                    _replace_or_remove(wrapper.module, original_node, replacement_node),
                    cst.Module,
                )
                patched_str = patched_module.code
                return LintPatch(0, CodePosition(1, 0), original_str, patched_str)

    def apply(self, original_module_str: str) -> str:
        return "".join(
            [
                original_module_str[: self.start_offset],
                self.patched_diff_str,
                original_module_str[self.start_offset + len(self.original_diff_str) :],
            ]
        )

    def minimize(self) -> "LintPatch":
        """
        Attempts to reduce the size of the patch. If you're going to apply multiple
        patches at once (without re-running the linter in between each patch), like
        `arc lint` does, you should attempt to minimize your patches to reduce the
        chance of overlapping changes.
        """
        return self._minimize_tail()._minimize_head()

    def _minimize_head(self) -> "LintPatch":
        """
        Strips any characters off of the head of the patch that are unchanged, reducing
        the size of the patch.
        """
        # line is 1-indexed, column is 0-indexed
        start_line, start_column = self.start_position.line, self.start_position.column

        matching_head_len = 0
        was_cr = False  # track \r\n newlines as a single newline
        for (a, b) in zip(self.original_diff_str, self.patched_diff_str):
            if a != b:
                break
            elif a == "\n":
                if not was_cr:
                    start_line += 1
                start_column = 0
                was_cr = False
            elif a == "\r":
                start_line += 1
                start_column = 0
                was_cr = True
            else:
                start_column += 1
                was_cr = False
            matching_head_len += 1

        return LintPatch(
            self.start_offset + matching_head_len,
            CodePosition(start_line, start_column),
            self.original_diff_str[matching_head_len:],
            self.patched_diff_str[matching_head_len:],
        )

    def _minimize_tail(self) -> "LintPatch":
        """
        Strips any characters off of the tail of the patch that are unchanged, reducing
        the size of the patch.
        """
        # this is easier than _minimize_head because we don't have to track line/column
        # numbers.
        matching_tail_len = 0
        for (a, b) in zip(
            reversed(self.original_diff_str), reversed(self.patched_diff_str)
        ):
            if a != b:
                break
            matching_tail_len += 1

        return LintPatch(
            self.start_offset,
            self.start_position,
            # We can't use `arr[:-offset]`` because when `offset` is `0`, we'd end up
            # with an empty string.
            self.original_diff_str[: len(self.original_diff_str) - matching_tail_len],
            self.patched_diff_str[: len(self.patched_diff_str) - matching_tail_len],
        )


def _replace_or_remove(
    parent: cst.CSTNode,
    original_node: cst.CSTNode,
    replacement_node: Union[cst.CSTNode, cst.FlattenSentinel, cst.RemovalSentinel],
) -> cst.CSTNode:
    if isinstance(replacement_node, (cst.RemovalSentinel, cst.FlattenSentinel)):
        return cst.ensure_type(parent.deep_remove(original_node), cst.CSTNode)
    else:
        return parent.deep_replace(original_node, replacement_node)
