# 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 typing import Any, Generator

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 (
    config_types,
    post_process,
    recipe_api,
    recipe_test_api,
)

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

PROPERTIES = pipeline_pb2.InputProperties


def RunSteps(
    api: recipe_api.RecipeScriptApi,
    props: pipeline_pb2.InputProperties,
):
    """Launch sequences of builders."""

    build_ids: list[int] = []

    for i, round_property in enumerate(props.rounds):
        with api.step.nest(f'round {i}') as pres:
            builders: list[str] = []
            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: dict[str, dict[str, Any]] = {
                '$pigweed/pipeline': {
                    'inside_a_pipeline': True,
                    'round': i,
                    'builds_from_previous_iteration': [
                        str(x) for x in build_ids
                    ],
                },
            }

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

            failures: list[str] = []
            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) -> Generator[recipe_test_api.TestData, None, None]:
    """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}')

    def drop():
        return api.post_process(post_process.DropExpectation)

    yield (
        api.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.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)
        + drop()
    )

    yield (
        api.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)
        + drop()
    )

    yield (
        api.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)
        + drop()
    )

    yield (
        api.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)
        + drop()
    )