# Copyright 2024 The Pigweed Authors
#
# Licensed under the Apache License, Version 2.0 (the "License"); you may not
# use this file except in compliance with the License. You may obtain a copy of
# the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
# License for the specific language governing permissions and limitations under
# the License.
"""Recipe module for updating dependencies in Bazel WORKSPACE or MODULE.bazel files.

This module provides the `BazelRollApi` class, which facilitates the automated
rolling of dependencies declared in Bazel build files. It supports updating:
  - `git_repository` (or `git_override`) rules by changing their `commit` attribute
    to the latest hash from a specified remote and branch.
  - `cipd_repository` rules by updating their `tag` attribute to the latest
    version resolved via `cipd_roll.select_version`.

Key features include:
  - Reading and parsing Bazel files (WORKSPACE, MODULE.bazel).
  - Locating specific dependency entries based on attributes like 'remote' (for Git)
    or 'path' (for CIPD).
  - Modifying the relevant version attribute (e.g., 'commit', 'tag').
  - Adding or updating metadata comments near the modified lines to indicate
    automatic updates.
  - Regenerating `MODULE.bazel.lock`.
  - Generating `Roll` objects for standardized commit messages.
  - Flexible matching of repository URLs and delayed file writes.
"""

from __future__ import annotations

import dataclasses
import re
from collections.abc import Callable, Sequence
from typing import TYPE_CHECKING

from recipe_engine import config_types, recipe_api

from PB.recipe_modules.pigweed.bazel_roll.cipd_package import BazelCipdPackage
from PB.recipe_modules.pigweed.bazel_roll.git_repository import GitRepository

from RECIPE_MODULES.fuchsia.git_roll_util import api as git_roll_util_api

if TYPE_CHECKING:  # pragma: no cover
    from RECIPE_MODULES.pigweed.checkout import api as checkout_api


def nwise(iterable, n):
    # nwise('ABCDEFG', 3) → ABC BCD CDE DEF EFG
    # See also
    # https://docs.python.org/3/library/itertools.html#itertools.pairwise
    iterator = iter(iterable)
    initial_items = [None]
    for _ in range(1, n):
        initial_items.append(next(iterator, None))
    items = tuple(initial_items)
    for x in iterator:
        items = (*items[1:], x)
        yield items


def trim_nearby_lines(lines: Sequence[LineProxy]) -> list[LineProxy]:
    # If there's a blank line in the second half, remove any following lines.
    for i in range(len(lines) // 2, len(lines)):
        text = lines[i].text or ''
        if not text.strip():
            lines = lines[0:i]
            break

    # If there's a blank line in the first half, remove any preceding lines.
    for i in range(len(lines) // 2, 0, -1):
        text = lines[i].text or ''
        if not text.strip():
            lines = lines[i:]
            break

    return lines


def proximity_sort_nearby_lines(
    lines: Sequence[LineProxy],
    center: int,
) -> list[LineProxy]:
    # Shift the order to be center-out instead of ascending.
    lines = [(abs(center - x.idx), x.idx, x) for x in lines]
    return [x[-1] for x in sorted(lines)]


@dataclasses.dataclass
class UpdateResult:
    """Holds the result of an attempt to update a value in a Bazel file.

    Attributes:
        old_value: The original value that was found and potentially replaced.
        project_name: The name of the project/repository associated with the
            update, extracted from attributes like 'name' or 'module_name'.
        finalize: An optional callable. If provided (e.g., when `delay_write`
            is True), this function must be called to actually perform the file
            modification.
    """

    old_value: str
    project_name: str | None
    finalize: Callable[[], None] | None = None


GIT_HASH_PATTERN = '[0-9a-fA-F]{40}'


class LineProxy:
    """A proxy object to interact with a line within a list of strings.

    This allows modifying a line in a list indirectly, which can be useful
    when passing around references to lines that might be replaced or
    deleted from the original list.
    """

    def __init__(self, lines: list[str | None], idx: int):
        """Initializes a LineProxy.

        Args:
            lines: The list of strings (lines of a file).
            idx: The index of the line this proxy represents.
        """
        self._lines = lines
        self._idx = idx

    @property
    def idx(self) -> int:
        """The index of the line in the original list."""
        return self._idx

    @property
    def text(self) -> str | None:
        """The text content of the line. Can be set to a new string or None."""
        return self._lines[self._idx]

    @text.setter
    def text(self, value: str | None) -> None:
        """Sets the text content of the line."""
        self._lines[self._idx] = value

    def __repr__(self) -> str:
        """String representation showing index and text."""
        return f'{self.idx} {self.text}'


def proxy(lines: list[str | None]) -> list[LineProxy]:
    """Creates a list of LineProxy objects from a list of strings."""
    return [LineProxy(lines, i) for i in range(len(lines))]


class BazelRollApi(recipe_api.RecipeApi):
    """API for updating dependencies in Bazel WORKSPACE or MODULE.bazel files.

    This includes updating `git_repository` commit hashes and `cipd_repository`
    tags. It provides utilities for reading these files, finding relevant
    dependency entries, and modifying their versions.
    """

    UpdateResult = UpdateResult
    GIT_HASH_PATTERN = GIT_HASH_PATTERN

    def __init__(self, *args, **kwargs) -> None:
        super().__init__(*args, **kwargs)
        self._read_cache: dict[config_types.Path, Sequence[str]] = {}

    def workspace_path(
        self,
        root: config_types.Path,
        value: str,
    ) -> config_types.Path:
        """Figure out the location of the WORKSPACE or MODULE.bazel file.

        If value is '', look for root / 'WORKSPACE' and root / 'MODULE.bazel'.
        If exactly one exists, return it. If not, error out.

        If root / value is a file, return it.

        If root / value is a directory, set root = root / value and apply the
        above logic. This enables applying this logic to subdirectories.

        Args:
            api: Recipe API object.
            root: Checkout root.
            value: Relative path specified in properties.

        Returns:
            Path to the WORKSPACE or MODULE.bazel file.
        """
        if value:
            value_path = root / value
            self.m.path.mock_add_file(value_path)
            if self.m.path.isfile(value_path):
                return value_path
            elif self.m.path.isdir(value_path):  # pragma: no cover
                root = value_path
            else:
                self.m.step.empty(  # pragma: no cover
                    f'{value_path} does not exist',
                    status='FAILURE',
                )

        workspace = root / 'WORKSPACE'
        module_bazel = root / 'MODULE.bazel'

        self.m.path.mock_add_file(workspace)

        if self.m.path.isfile(module_bazel) and self.m.path.isfile(workspace):
            self.m.step.empty(  # pragma: no cover
                f'{module_bazel} and {workspace} both exist',
                status='FAILURE',
            )

        if self.m.path.isfile(module_bazel):
            return module_bazel  # pragma: no cover

        if self.m.path.isfile(workspace):
            return workspace

        self.m.step.empty(  # pragma: no cover
            'no WORKSPACE or MODULE.bazel file found',
            status='FAILURE',
        )

    def update_cipd_package(
        self,
        checkout: checkout_api.CheckoutContext,
        package: BazelCipdPackage,
        platforms: Sequence[str] | None = None,
    ) -> list[git_roll_util_api.Roll]:
        """Updates a `cipd_repository` entry in a Bazel workspace file.

        This method finds a `cipd_repository` rule in the specified
        WORKSPACE or MODULE.bazel file that matches the given package spec.
        It then determines the new version tag for this CIPD package (using
        `cipd_roll.select_version`) and updates the `tag` attribute in the
        Bazel file.

        Args:
            checkout: The checkout context.
            package: A protobuf message defining the CIPD package to update,
                including its name, spec, workspace path, ref, and tag key.
            platforms: An optional list of platforms to consider for the CIPD
                package. Defaults to a standard set of Mac and Linux platforms.

        Returns:
            A list containing a `Roll` object (from `cipd_roll`) if the
            package was updated, or an empty list if no update was needed.
        """
        if not package.name:
            package.name = self.m.cipd_roll.package_name(package.spec)

        # TODO: b/234874582 - Support Windows.
        if platforms is None:
            platforms = ('mac-amd64', 'mac-arm64', 'linux-amd64')

        with self.m.step.nest(package.name):
            workspace_path = self.workspace_path(
                checkout.root,
                package.workspace_path,
            )

            attrs = self.retrieve_attributes(
                project_value=package.spec,
                path=workspace_path,
                search_key='path',
            )
            assert len(attrs) == 1, repr(attrs)

            old_version = attrs[0]['tag']

            new_version = self.m.cipd_roll.select_version(
                spec=package.spec,
                ref=package.ref,
                tag=package.tag,
                allow_mismatched_refs=package.allow_mismatched_refs,
                platforms=platforms,
                old_version=old_version,
            )
            if not new_version:
                # select_version() failing is adequately covered by the
                # cipd_roll module.
                return []  # pragma: no cover

            update_result = self.update_value(
                project_value=package.spec,
                new_value=new_version,
                path=workspace_path,
                delay_write=True,
                search_key='path',
                value_key='tag',
                value_pattern=f'{package.tag}:.+',
                matches=checkout.remotes_equivalent,
            )
            if not update_result:
                # update_value() failing is adequately covered by
                # update_git_repository().
                self.m.step.empty(  # pragma: no cover
                    'failed to update commit hash',
                    status='FAILURE',
                )

            roll = self.m.cipd_roll.Roll(
                package_name=package.name,
                package_spec=package.spec,
                old_version=old_version,
                new_version=new_version,
            )

            update_result.finalize()

            return [roll]

    def update_git_repository(
        self,
        checkout: checkout_api.CheckoutContext,
        git_repository: GitRepository,
    ) -> list[git_roll_util_api.Roll]:
        """Updates a `git_repository` or `git_override` entry in a Bazel file.

        This method finds a `git_repository` (or `git_override`) rule in the
        specified WORKSPACE or MODULE.bazel file that matches the given remote URL.
        It then resolves the latest commit hash for the repository's specified
        branch and updates the `commit` attribute in the Bazel file.

        Args:
            checkout: The checkout context.
            git_repository: A protobuf message defining the Git repository to
                update, including its remote URL, branch, workspace path, and
                name.

        Returns:
            A list containing a `Roll` object (from `git_roll_util`) if the
            repository was updated, or an empty list if no update was needed
            (e.g., already up-to-date or a backwards roll).

        Raises:
            StepFailure: If the `update_value` call fails to find or update the
                commit hash.
        """
        project_name = git_repository.name
        if not project_name:
            project_name = git_repository.remote.split('/')[-1]
            project_name = project_name.removesuffix('.git')

        with self.m.step.nest(project_name):
            workspace_path = self.workspace_path(
                checkout.root,
                git_repository.workspace_path,
            )
            git_repository.branch = git_repository.branch or 'main'

            new_revision = self.m.git_roll_util.resolve_new_revision(
                git_repository.remote,
                git_repository.branch,
                checkout.remotes_equivalent,
            )

            update_result = self.update_value(
                project_value=git_repository.remote,
                new_value=new_revision,
                path=workspace_path,
                delay_write=True,
                search_key='remote',
                value_key='commit',
                value_pattern=GIT_HASH_PATTERN,
                matches=checkout.remotes_equivalent,
            )
            if not update_result:
                self.m.step.empty(
                    'failed to update commit hash',
                    status='FAILURE',
                )

            try:
                roll = self.m.git_roll_util.get_roll(
                    repo_url=git_repository.remote,
                    repo_short_name=project_name,
                    old_rev=update_result.old_value,
                    new_rev=new_revision,
                )

            except self.m.git_roll_util.BackwardsRollError:
                props = self.m.step.empty('output property').presentation
                props.properties[project_name] = {
                    'remote': git_repository.remote,
                    'revision': new_revision,
                }

                return []

            update_result.finalize()

            return [roll]

    def _read(
        self,
        path: config_types.Path,
        num_nearby_lines: int,
        allow_cache: bool = True,
    ) -> list[str | None]:
        """Reads a file and pads it with empty lines for context searching.

        The file content is cached in `self._read_cache` to avoid redundant
        reads. `num_nearby_lines` empty strings are added to the beginning and
        end of the actual file lines. This padding simplifies searching for
        patterns near the start or end of the file by ensuring there's always
        enough "context" lines available for `nwise` iteration.

        Args:
            path: The path to the file to read.
            num_nearby_lines: The number of empty lines to pad at the beginning
                and end.
            allow_cache: If True, uses the cached content if available.

        Returns:
            A list of strings representing the lines of the file, padded with
            empty strings.
        """
        lines: list[str | None] = [''] * num_nearby_lines

        if allow_cache and path in self._read_cache:
            lines.extend(self._read_cache[path])
        else:
            self._read_cache[path] = (
                self.m.file.read_text(
                    f'read old {path.name}',
                    path,
                    test_data=self.m.bazel_roll.test_api.TEST_WORKSPACE_FILE,
                )
                .strip()
                .splitlines()
            )
            lines.extend(self._read_cache[path])

        lines.extend([''] * num_nearby_lines)
        return lines

    def _get_matching_groups(
        self,
        lines: Sequence[str],
        num_nearby_lines: int,
        project_value: str,
        search_key: str,
        replace: bool = False,
        matches: Callable[[str, str], bool] | None = None,
    ) -> list[tuple[LineProxy, list[LineProxy]]]:
        """Get a list of "groups" of lines from a Bazel file.

        Example for search_key "abc" and matches "def:.+". These would be
        ignored:

            foo(
                abc = "def",
                foo = "bar",
                uvw = "xyz",
            )

            foo(
                abc = "abc",
                foo = "foo",
                xyz = "xyz",
            )

        And this would match:

            foo(
                abc = "def:123",
                foo = "foo",
                xyz = "xyz",
            )

        Retrieve the matching groups.
        """
        matching_groups: list[tuple[LineProxy, list[LineProxy]]] = []

        # This is more constraining than necessary, but ensures there aren't
        # special characters that confuse the regex.
        assert search_key.isidentifier(), search_key

        for nearby_lines in nwise(proxy(lines), num_nearby_lines * 2 + 1):
            curr = nearby_lines[len(nearby_lines) // 2]
            match = re.search(
                rf'^\s*{search_key}\s*=\s*"(?P<search>[^"]+)",?\s*$',
                curr.text,
            )
            if not match:
                continue

            match_value = match.group('search')

            def equals(x, y):
                # TODO: b/376915917 - Remove no cover pragma.
                return x == y  # pragma: no cover

            if not matches:
                # TODO: b/376915917 - Remove no cover pragma.
                matches = equals  # pragma: no cover

            if matches(match_value, project_value):
                pres = self.m.step.empty(
                    f'found matching {search_key} {match_value!r}'
                ).presentation
                pres.logs['lines'] = [repr(x) for x in nearby_lines]
                matching_groups.append((curr, nearby_lines))

                if replace and match_value != project_value:
                    curr.text = curr.text.replace(match_value, project_value)

            else:
                pres = self.m.step.empty(
                    f'found other {search_key} {match_value!r}'
                ).presentation
                pres.logs['lines'] = [repr(x) for x in nearby_lines]

        return matching_groups

    def _process_nearby_lines(
        self,
        nearby_lines: Sequence[LineProxy],
    ) -> list[LineProxy]:
        """Processes a list of nearby lines for attribute extraction.

        This involves:
        1. Trimming blank lines from the periphery of the `nearby_lines` list.
        2. Sorting the remaining lines by their proximity to the original
           center line (the line that matched the `search_key`).

        This helps in prioritizing lines closer to the `search_key` match when
        looking for other attributes like `value_key` or `title_keys`.

        Args:
            nearby_lines: A sequence of `LineProxy` objects representing the
                context around a matched line.

        Returns:
            A processed list of `LineProxy` objects, trimmed and sorted.
        """
        pres = self.m.step.empty('lines').presentation
        center = nearby_lines[len(nearby_lines) // 2].idx
        pres.logs['0_center'] = [str(center)]

        pres.logs['1_orig'] = [repr(x) for x in nearby_lines]

        nearby_lines = trim_nearby_lines(nearby_lines)
        pres.logs['2_trimmed'] = [repr(x) for x in nearby_lines]

        nearby_lines = proximity_sort_nearby_lines(nearby_lines, center)
        pres.logs['3_sorted'] = [repr(x) for x in nearby_lines]

        return nearby_lines

    def retrieve_attributes(
        self,
        project_value: str,
        path: config_types.Path,
        search_key: str,
        num_nearby_lines: int = 10,
        title_keys: Sequence[str] = ('name', 'module_name'),
        matches: Callable[[str, str], bool] | None = None,
    ) -> list[dict[str, str]]:
        """Retrieves attributes from matching groups in a Bazel file.

        This function searches a Bazel file (WORKSPACE or MODULE.bazel) for
        blocks of attributes (like `git_repository` or `cipd_repository`
        rules). It identifies these blocks by a `search_key` (e.g., 'remote'
        or 'path') matching a `project_value`. For each matching block, it
        extracts all key-value attributes.

        Args:
            project_value: The value to match for the `search_key` (e.g., a
                specific remote URL or CIPD package path).
            path: The path to the Bazel file.
            search_key: The attribute key used to identify relevant blocks
                (e.g., 'remote', 'path').
            num_nearby_lines: The number of lines around a `search_key` match to
                consider as part of the attribute block.
            title_keys: A sequence of attribute keys (e.g., 'name',
                'module_name') that, if found, will also populate an 'name'
                field in the returned dictionary for that block.
            matches: An optional callable `(found_value, project_value) -> bool`
                to determine if a `search_key`'s value matches the target
                `project_value`. Defaults to exact string equality.

        Returns:
            A list of dictionaries, where each dictionary represents a matched
            attribute block. Keys are attribute names, and values are their
            corresponding string values. An additional 'name' key might be
            present if one of the `title_keys` was found.
        """
        lines = self._read(path, num_nearby_lines)

        assert isinstance(title_keys, (list, tuple)), repr(title_keys)

        matching_groups = self._get_matching_groups(
            lines=lines,
            num_nearby_lines=num_nearby_lines,
            project_value=project_value,
            search_key=search_key,
            matches=matches,
        )

        results: list[dict[str, str]] = []
        for _, nearby_lines in matching_groups:
            entry = {}
            nearby_lines = self._process_nearby_lines(nearby_lines)

            for line in nearby_lines:
                if match := re.search(
                    r'^\s*([\w_]+)\s*=\s*(\S.*),\s*$',
                    line.text,
                ):
                    entry[match.group(1)] = match.group(2).strip('"')
                    if match.group(1) in title_keys:
                        entry['name'] = match.group(2).strip('"')

            if entry:
                results.append(entry)

        return results

    def update_value(
        self,
        *,
        project_value: str,
        new_value: str,
        path: config_types.Path,
        search_key: str,
        value_key: str,
        num_nearby_lines: int = 10,
        replace: bool = False,
        delay_write: bool = False,
        title_keys: Sequence[str] = ('name', 'module_name'),
        value_pattern: str | None = None,
        matches: Callable[[str, str], bool] | None = None,
    ) -> UpdateResult | None:
        """Updates a specific attribute's value within a matching block in a Bazel file.

        This function locates a block of attributes (e.g., a `git_repository`
        rule) using `project_value` and `search_key`. Within that block, it
        finds the line corresponding to `value_key` (e.g., 'commit' or 'tag')
        and updates its value to `new_value`. It also adds/updates metadata
        comments near the modified line.

        Args:
            project_value: The value for the `search_key` to identify the block.
            new_value: The new value to set for the `value_key` attribute.
            path: The path to the Bazel file.
            search_key: The attribute key to identify the block (e.g., 'remote').
            value_key: The attribute key whose value needs to be updated (e.g., 'commit').
            num_nearby_lines: Context lines to consider around matches.
            replace: If True, and `search_key`'s value doesn't exactly match
                     `project_value` but `matches()` returns True, the
                     `search_key`'s value will be updated to `project_value`.
            delay_write: If True, the file modification is not performed
                         immediately. Instead, a `finalize` callable is returned
                         in `UpdateResult`, which must be called to write changes.
            title_keys: Keys to identify the project name for `UpdateResult`.
            value_pattern: A regex pattern to match the existing value of `value_key`.
                           Defaults to a generic pattern if None.
            matches: A callable for flexible matching of `project_value`.

        Returns:
            An `UpdateResult` object containing the old value and project name,
            and potentially a `finalize` function if `delay_write` is True.
            Returns None if the `project_value` or `value_key` is not found.
        """
        lines = self._read(path, num_nearby_lines)

        matching_groups = self._get_matching_groups(
            lines=lines,
            num_nearby_lines=num_nearby_lines,
            project_value=project_value,
            replace=replace,
            search_key=search_key,
            matches=matches,
        )

        if not matching_groups:
            self.m.step.empty(
                f'could not find {search_key} {project_value} in {path}',
            )
            return None

        project_names: list[str] = []
        if value_pattern is None:
            # TODO: b/376915917 - Remove no cover pragma.
            value_pattern = '[^"]+'  # pragma: no cover

        for matching_line, nearby_lines in matching_groups:
            nearby_lines = self._process_nearby_lines(nearby_lines)

            value_rx = re.compile(
                rf'^(?P<prefix>\s*{value_key}\s*=\s*")'
                rf'(?P<value>{value_pattern})'
                r'(?P<suffix>",?\s*)$'
            )

            for line in nearby_lines:
                if match := value_rx.search(line.text):
                    idx = line.idx
                    break
            else:
                pres = self.m.step.empty(
                    f'could not find {value_key} line adjacent to '
                    f'{matching_line.text!r} in {path}',
                ).presentation
                pres.step_summary_text = repr(nearby_lines)
                return None

            old_value = match.group('value')

            prefix = match.group('prefix')
            suffix = match.group('suffix')
            lines[idx] = f'{prefix}{new_value}{suffix}'

            # Remove all existing metadata lines in this git_repository() entry.
            idx2 = idx - 1
            while lines[idx2].strip().startswith('# ROLL: '):
                lines[idx2] = None
                idx2 -= 1

            ws_prefix = re.search(r'^\s*', prefix).group(0)
            comment_prefix = f'{ws_prefix}# ROLL: '

            comment_lines = (
                f'{comment_prefix}Warning: this entry is automatically '
                'updated by',
                f'{comment_prefix}{self.m.buildbucket.builder_url()}',
            )

            lines[idx] = '\n'.join(comment_lines + (lines[idx],))

            assert isinstance(title_keys, (list, tuple)), repr(title_keys)

            # This is more restrictive than necessary but helps keep the regex
            # simple.
            for key in title_keys:
                assert isinstance(key, str), repr(key)
                assert key.isidentifier(), repr(key)

            for line in nearby_lines:
                keys = '|'.join(title_keys)
                if match := re.search(
                    rf'^\s*(?:{keys})\s*=\s*"(?P<value>[^"]+)",?\s*$',
                    line.text or '',
                ):
                    project_names.append(match.group('value'))
                    break

        def write():
            self.m.file.write_text(
                f'write new {path.name}',
                path,
                ''.join(
                    f'{x}\n'
                    for x in lines[num_nearby_lines:-num_nearby_lines]
                    if x is not None
                ),
            )

        if delay_write:
            return UpdateResult(
                old_value=old_value,
                project_name=', '.join(project_names),
                finalize=write,
            )

        else:
            write()

            return UpdateResult(
                old_value=old_value,
                project_name=', '.join(project_names),
            )