# Copyright 2025 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.
"""Split combined triggers into separate builds by branch."""

from __future__ import annotations

import dataclasses
import urllib.parse
from typing import Any

from recipe_engine import recipe_api

from PB.go.chromium.org.luci.buildbucket.proto import (
    builds_service as builds_service_pb,
)
from PB.go.chromium.org.luci.buildbucket.proto import common as common_pb
from PB.go.chromium.org.luci.scheduler.api.scheduler.v1 import (
    triggers as triggers_pb,
)
from PB.recipe_engine import result as result_pb
from PB.recipe_modules.pigweed.split_triggers.properties import InputProperties

USER_AGENT = 'user_agent'
SPLIT_TRIGGERS = 'split_triggers'


@dataclasses.dataclass(frozen=True)
class RemoteBranch:
    remote: str
    ref: str


class SplitTriggersApi(recipe_api.RecipeApi):
    """Split combined triggers into separate builds by branch."""

    def __init__(self, props: InputProperties, *args: Any, **kwargs: Any):
        super().__init__(*args, **kwargs)
        self._props = props

    def __call__(self) -> result_pb.RawResult | None:
        """Split combined triggers into separate builds by branch."""
        if self.m.buildbucket_util.is_tryjob:
            return

        # Make sure we never end up in infinite recursion.
        for tag in self.m.buildbucket.build.tags:
            if tag.key == USER_AGENT and tag.value == SPLIT_TRIGGERS:
                return None

        with self.m.step.nest('split triggers'):
            return self._split()

    def _split(self) -> result_pb.RawResult | None:
        pres = self.m.step.empty('properties').presentation
        pres.step_summary_text = repr(self.m.properties.thaw())

        with self.m.step.nest('processing triggers'):
            triggers_by_branch: dict[
                RemoteBranch, list[triggers_pb.Trigger]
            ] = {}
            for trigger in self.m.scheduler.triggers:
                if trigger.WhichOneof('payload') != 'gitiles':
                    self.m.step.empty('one non-gitiles trigger')
                    continue

                key = RemoteBranch(trigger.gitiles.repo, trigger.gitiles.ref)
                self.m.step.empty(
                    f'one trigger for {key.remote} {key.ref} '
                    f'{trigger.gitiles.revision}',
                )
                triggers_by_branch.setdefault(key, [])
                triggers_by_branch[key].append(trigger)

        requests: list[builds_service_pb.ScheduleBuildRequest] = []

        schedule_kwargs = {}
        schedule_kwargs['builder'] = self.m.buildbucket.build.builder.builder
        schedule_kwargs['inherit_buildsets'] = False
        schedule_kwargs['experiments'] = {'pigweed.disable_git_cache': True}
        schedule_kwargs['can_outlive_parent'] = True
        schedule_kwargs['tags'] = [
            common_pb.StringPair(key=USER_AGENT, value=SPLIT_TRIGGERS),
        ]

        num_branches = len(triggers_by_branch)

        if num_branches == 0:
            pres = self.m.step.empty('no branches').presentation
            pres.properties['split triggers'] = {
                'to be split': False,
                'reason': 'no branches',
                'enabled': self._props.enabled,
            }
            return None

        if num_branches == 1:
            pres = self.m.step.empty('only one branch').presentation
            pres.properties['split triggers'] = {
                'to be split': False,
                'reason': 'only one branch',
                'enabled': self._props.enabled,
            }
            return None

        with self.m.step.nest('filtering'):
            for triggers in triggers_by_branch.values():
                # The most recent commit on a branch shows up last in the
                # list of commits, so we only need to trigger the build on
                # it.
                trigger = triggers[-1].gitiles
                parsed = urllib.parse.urlparse(trigger.repo)

                commit = common_pb.GitilesCommit(
                    host=parsed.hostname,
                    project=parsed.path.lstrip('/'),
                    id=trigger.revision,
                    ref=trigger.ref,
                )

                requests.append(
                    self.m.buildbucket.schedule_request(
                        gitiles_commit=commit,
                        **schedule_kwargs,
                    )
                )
                self.m.step.empty(
                    f'selected {trigger.repo} {trigger.ref} {trigger.revision}',
                )

        with self.m.step.nest(f'launching {len(requests)} builds') as pres:
            if self._props.enabled:
                summary = [
                    f'Split the triggers into {len(requests)} based on trigger '
                    'branches',
                ]
                builds = self.m.buildbucket.schedule(requests)

                pres.properties['split triggers'] = {
                    'to be split': True,
                    'reason': f'{len(requests)} branches',
                    'enabled': self._props.enabled,
                }

                for build in builds:
                    branch = build.input.gitiles_commit.ref.removeprefix(
                        'refs/heads/',
                    )
                    summary.append(
                        f'* {branch}: {self.m.buildbucket.build_url(build=build)}',
                    )

                return result_pb.RawResult(
                    summary_markdown='\n'.join(summary),
                    status=common_pb.SUCCESS,
                )

            else:
                self.m.step.empty('dry-run')

                pres.properties['split triggers'] = {
                    'to be split': True,
                    'reason': f'{len(requests)} branches',
                    'enabled': self._props.enabled,
                }

                for request in requests:
                    branch = request.gitiles_commit.ref.removeprefix(
                        'refs/heads/',
                    )
                    self.m.step.empty(branch)

                return None