# Copyright 2022 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 for launching sequences of builders."""

import re

from PB.go.chromium.org.luci.buildbucket.proto import common as common_pb2
import PB.recipes.pigweed.pipeline as pipeline_pb2
from PB.recipe_engine import result
from recipe_engine import post_process

DEPS = [
    'fuchsia/status_check',
    'fuchsia/subbuild',
    'recipe_engine/properties',
    'recipe_engine/step',
]

PROPERTIES = pipeline_pb2.InputProperties


def RunSteps(api, props):
    """Launch sequences of builders."""

    build_ids = []

    for i, round_property in enumerate(props.rounds):
        with api.step.nest(f'round {i}') as pres:
            builders = []
            for builder_property in round_property.builders:
                # TODO(b/245788264) Support non-current project.
                assert not builder_property.project
                # TODO(b/245788264) Support non-current bucket.
                assert not builder_property.bucket
                builders.append(builder_property.builder)

            extra_props = {
                '$pigweed/pipeline': {
                    'inside_a_pipeline': True,
                    'round': i,
                    'builds_from_previous_iteration': build_ids,
                },
            }

            launched_builds = api.subbuild.launch(
                builders, pres, extra_properties=extra_props
            )
            for builder, build in launched_builds.items():
                pres.links[builder] = build.url
            build_ids = [x.build_id for x in launched_builds.values()]
            collected_builds = api.subbuild.collect(build_ids)

            failures = []
            with api.step.defer_results():
                for build in collected_builds.values():
                    with api.step.nest(build.builder) as sub_pres:
                        if build.build_proto.status != common_pb2.SUCCESS:
                            sub_pres.status = 'FAILURE'
                            failures.append(build.builder)

            if failures:
                return result.RawResult(
                    summary_markdown=f'Failed builds: {", ".join(failures)}',
                    status=common_pb2.FAILURE,
                )


def GenTests(api):  # pylint: disable=invalid-name
    """Create tests."""

    def builder(name):
        match = re.match(
            r'(?:(?:(?P<project>[^/]+)/)?(?P<bucket>[^/]+)/)?(?P<name>[^/]+)$',
            name,
        )
        assert match
        return pipeline_pb2.Builder(
            project=match.group('project') or '',
            bucket=match.group('bucket') or '',
            builder=match.group('name'),
        )

    def round(*builders, timeout_sec=0):
        return pipeline_pb2.Round(
            builders=list(builders), timeout_sec=timeout_sec,
        )

    def properties(*rounds):
        return api.properties(pipeline_pb2.InputProperties(rounds=list(rounds)))

    def result_ci(name, id, status='SUCCESS', **kwargs):
        return api.subbuild.ci_build_message(
            builder=name, build_id=id, status=status, **kwargs,
        )

    def assert_step_success(round, name):
        return api.post_process(
            post_process.StepSuccess, f'round {round}.{name}'
        )

    def assert_step_failure(round, name):
        return api.post_process(
            post_process.StepFailure, f'round {round}.{name}'
        )

    def assert_num_rounds(round):
        return api.post_process(
            post_process.MustRun, f'round {round-1}'
        ) + api.post_process(post_process.DoesNotRun, f'round {round}')

    yield (
        api.status_check.test('one_builder')
        + properties(round(builder('a')))
        + api.subbuild.child_build_steps(
            [result_ci('a', 1000)],
            launch_step='round 0',
            collect_step='round 0',
        )
        + assert_step_success(0, 'a')
        + assert_num_rounds(1)
    )

    yield (
        api.status_check.test('two_builders')
        + properties(round(builder('a'), builder('b')),)
        + api.subbuild.child_build_steps(
            [result_ci('a', 1000), result_ci('b', 1001)],
            launch_step='round 0',
            collect_step='round 0',
        )
        + assert_step_success(0, 'a')
        + assert_step_success(0, 'b')
        + assert_num_rounds(1)
    )

    yield (
        api.status_check.test('two_rounds')
        + properties(round(builder('a')), round(builder('b')),)
        + api.subbuild.child_build_steps(
            [result_ci('a', 1000)],
            launch_step='round 0',
            collect_step='round 0',
        )
        + api.subbuild.child_build_steps(
            [result_ci('b', 2000)],
            launch_step='round 1',
            collect_step='round 1',
        )
        + assert_step_success(0, 'a')
        + assert_step_success(1, 'b')
        + assert_num_rounds(2)
    )

    yield (
        api.status_check.test('first_round_failure', status='failure')
        + properties(round(builder('a'), builder('b')), round(builder('c')),)
        + api.subbuild.child_build_steps(
            [result_ci('a', 1000), result_ci('b', 1001, status='FAILURE')],
            launch_step='round 0',
            collect_step='round 0',
        )
        + assert_step_success(0, 'a')
        + assert_step_failure(0, 'b')
        + assert_num_rounds(1)
    )

    yield (
        api.status_check.test('second_round_failure', status='failure')
        + properties(
            round(builder('a'), builder('b')),
            round(builder('c')),
            round(builder('d')),
        )
        + api.subbuild.child_build_steps(
            [result_ci('a', 1000), result_ci('b', 1001)],
            launch_step='round 0',
            collect_step='round 0',
        )
        + api.subbuild.child_build_steps(
            [result_ci('c', 2000, 'FAILURE')],
            launch_step='round 1',
            collect_step='round 1',
        )
        + assert_step_success(0, 'a')
        + assert_step_success(0, 'b')
        + assert_step_failure(1, 'c')
        + assert_num_rounds(2)
    )