# 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.

import attrs
import itertools
from typing import Any

from google.protobuf import json_format
from PB.go.chromium.org.luci.buildbucket.proto import (
    build as build_pb2,
    builds_service as builds_service_pb2,
)
from PB.recipes.pigweed.multi_roller import Builder, InputProperties

DEPS = [
    "fuchsia/buildbucket_util",
    "fuchsia/gerrit",
    "fuchsia/subbuild",
    "recipe_engine/buildbucket",
    "recipe_engine/json",
    "recipe_engine/properties",
    "recipe_engine/step",
    "recipe_engine/time",
]

PROPERTIES = InputProperties

LABEL_STATUSES = ("rejected", "approved", "disliked", "recommended")
PRESUBMIT_VERIFIED_LABEL = "Presubmit-Verified"

DEFAULT_POLL_INTERVAL_SECONDS = 5 * 60


@attrs.define
class GerritChange:
    host: str
    change_id: str

    @property
    def name(self):
        return f'{self.host}:{self.change_id}'

    @property
    def url(self):
        return f'https://{self.host}-review.googlesource.com/q/{self.change_id}'


def RunSteps(api, props):
    """The multi_roller recipe launches and synchronizes multiple rollers.

    The rollers to launch is specified via the input properties. The recipe
    will use the subbuild recipe_module to launch the rollers. Once each roller
    has created a roll CL, it will also output a property that exposes the
    gerrit host and change-id of that roll CL. The rollers will then wait for a
    Rolls-Synced label. Multi_roller watches over all the rollers and only sets
    the Rolls-Synced label on each roll CL once all of the roll CLs pass
    presubmit checks.
    """
    # Launch each sub-roller.
    with api.step.nest("Launching Rollers") as pres:
        builders: list[str] = []
        for builder_property in props.rollers_to_launch:
            # 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)

        # Launch all sub-rollers with the same gerrit topic.
        # Append the multi_roller's build id to the topic to make it unique.
        launched_rollers: dict[str, api.subbuild.SubbuildResult] = (
            api.subbuild.launch(
                builders,
                pres,
                extra_properties={
                    "override_auto_roller_options": {
                        "push_options": (
                            f"topic=pigweed-s5400-roller-{api.buildbucket.build.id}"
                        )
                    }
                },
            )
        )
        roller_build_ids: list[int] = [
            int(x.build_id) for x in launched_rollers.values()
        ]

        # Calculate how often to poll gerrit and buildbucket.
        props.poll_interval_secs = (
            props.poll_interval_secs or DEFAULT_POLL_INTERVAL_SECONDS
        )

        # Wait for sub-rollers to output the change ID of their roll CLs.
        roller_changes = get_roller_changes(api, props, roller_build_ids)

        # Wait for roll CLs to pass presubmits.
        presubmits_passed = wait_for_presubmit(api, props, roller_changes)

        # Add Rolls-Synced label value on roll CLs.
        roll_sync_label_value = 1 if presubmits_passed else -1
        for change in roller_changes:
            api.gerrit.set_review(
                name=f"Set {props.rolls_synced_label} Label on {change.change_id}",
                change_id=change.change_id,
                labels={props.rolls_synced_label: roll_sync_label_value},
                host=change.host,
                test_data=api.json.test_api.output({}),
            )


def _build_url(build_id):
    return f"https://ci.chromium.org/ui/b/{build_id}"


def get_roller_changes(api, props, build_ids: list[int]) -> list[GerritChange]:
    """Uses buildbucket api to check the output property of builds.

    All sub-roller builds will output an output property that contains the host
    and change ID of their roll CL.

    Returns
        A list of GerritChange objects corresponding to the sub-roller CLs.
    """
    roller_changes: dict[str, list[GerritChange]] = {}
    remaining_build_ids: list[str] = list(build_ids)
    for i in itertools.count():
        with api.step.nest(f"Get sub-roller change IDs ({i})") as pres:
            roller_builds = api.buildbucket.get_multi(remaining_build_ids)
            for build_id, build in roller_builds.items():
                output_props = json_format.MessageToDict(
                    build.output.properties
                )
                if "gerrit_changes" in output_props:
                    with api.step.nest(
                        f"got changes for {build_id}"
                    ) as change_pres:
                        change_pres.links[build_id] = _build_url(build_id)
                        for chg in output_props["gerrit_changes"]:
                            change = GerritChange(
                                host=chg["host"], change_id=chg["change_id"]
                            )
                            change_pres.links[change.name] = change.url
                            roller_changes.setdefault(build_id, [])
                            roller_changes[build_id].append(change)
                        remaining_build_ids.remove(build_id)

            if not remaining_build_ids:
                result = list(
                    itertools.chain.from_iterable(roller_changes.values())
                )
                pres.step_summary_text = f"Got all change ids: {result}"
                return result

            if _on_last_iteration(api, props.poll_interval_secs):
                break
            api.time.sleep(props.poll_interval_secs)

    with api.step.nest("Failed to retrieve some change ids") as pres:
        for build_id in remaining_build_ids:
            pres.links[build_id] = _build_url(build_id)
        pres.status = "FAILURE"
        raise api.step.StepFailure("Failed to retrieve some change ids")


def wait_for_presubmit(api, props, roller_changes: list[GerritChange]):
    """Wait for all sub-roller roll CLs to pass presubmit checks.

    Returns:
        Whether or not all sub-roller CLs passed presubmits.
    """
    success_status = {"approved"}
    failure_status = set(LABEL_STATUSES) - success_status
    remaining_rollers: list[GerritChange] = list(roller_changes)
    for i in itertools.count():
        with api.step.nest(
            f"Wait for {PRESUBMIT_VERIFIED_LABEL} for roll CLs ({i})"
        ) as pres:
            pres.step_summary_text = f"Waiting on {PRESUBMIT_VERIFIED_LABEL} for: {remaining_rollers}"
            for change in list(remaining_rollers):
                step = api.gerrit.change_details(
                    f"Get CL details for CL {change.change_id}",
                    change.change_id,
                    host=change.host,
                )
                details = step.json.output

                # If the roll CL was manually abandoned or merged, stop waiting.
                if details["status"] == "ABANDONED":
                    api.step.empty(
                        f"Roll CL {change.change_id} was manually abandoned."
                    )
                    return False
                if details["status"] == "MERGED":
                    api.step.empty(
                        f"Roll CL {change.change_id} was manually merged."
                    )
                    remaining_rollers.remove(change)
                    continue

                with api.step.nest(
                    f"Check CL {change.change_id}"
                ) as label_pres:
                    label_pres.step_summary_text = (
                        f"{PRESUBMIT_VERIFIED_LABEL} not set"
                    )
                    if PRESUBMIT_VERIFIED_LABEL not in details["labels"]:
                        label_pres.step_summary_text = (
                            "f{PRESUBMIT_VERIFIED_LABEL} does not exist"
                        )
                        return False

                    if failure_status.intersection(
                        details["labels"][PRESUBMIT_VERIFIED_LABEL]
                    ):
                        label_pres.step_summary_text = "Presubmit checks failed"
                        return False

                    if success_status.intersection(
                        details["labels"][PRESUBMIT_VERIFIED_LABEL]
                    ):
                        label_pres.step_summary_text = "Presubmit checks passed"
                        remaining_rollers.remove(change)

            if not remaining_rollers:
                pres.step_summary_text = (
                    f"All roll CLs have {PRESUBMIT_VERIFIED_LABEL}"
                )
                return True

            # Don't sleep after the final check.
            if _on_last_iteration(api, props.poll_interval_secs):
                return False
            api.time.sleep(props.poll_interval_secs)


def _on_last_iteration(api, poll_interval_secs):
    return _secs_remaining_in_build(api) < poll_interval_secs * 3 // 2


def _secs_remaining_in_build(api):
    current_time = api.time.time()
    elapsed_time = current_time - api.buildbucket.build.start_time.seconds
    return api.buildbucket.build.execution_timeout.seconds - elapsed_time


def GenTests(api):
    def test(*args, **kwargs):
        kwargs.setdefault("execution_timeout", 100)
        return (
            api.buildbucket_util.test(*args, **kwargs)
            + api.time.seed(0)
            + api.time.step(1)
        )

    def builder(name):
        return Builder(project="", bucket="", builder=name)

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

    def properties(*builders, **kwargs):
        props = InputProperties(**kwargs)
        props.rollers_to_launch.extend(list(builders))
        props.poll_interval_secs = 5
        props.rolls_synced_label = "Rolls-Synced"
        return api.properties(props)

    def simulate_subroller_launch(builds, launch_step="Launching Rollers"):
        responses = []
        for build in builds:
            responses.append(
                dict(schedule_build=dict(id=build.id, builder=build.builder))
            )
        return api.buildbucket.simulated_schedule_output(
            step_name=f"{launch_step}.schedule",
            batch_response=builds_service_pb2.BatchResponse(
                responses=responses
            ),
        )

    def simulate_get_multi(builds, launch_step):
        return api.buildbucket.simulated_get_multi(
            builds=builds,
            step_name=f"{launch_step}.buildbucket.get_multi",
        )

    yield (
        test("success with two sub-rollers")
        + properties(builder("roller-1"), builder("roller-2"))
        + simulate_subroller_launch(
            builds=[
                result_ci(name="roller-1", id=1000),
                result_ci(name="roller-2", id=1001),
            ],
        )
        + simulate_get_multi(
            builds=[
                result_ci(
                    name="roller-1",
                    id=1000,
                    output_props={
                        "gerrit_changes": [
                            {"host": "change_host_1", "change_id": "11111111"}
                        ]
                    },
                ),
                result_ci(
                    name="roller-2",
                    id=1001,
                    # output property not set yet!
                    output_props={},
                ),
            ],
            launch_step="Launching Rollers.Get sub-roller change IDs (0)",
        )
        + simulate_get_multi(
            builds=[
                result_ci(
                    name="roller-2",
                    id=1001,
                    output_props={
                        "gerrit_changes": [
                            {"host": "change_host_2", "change_id": "22222222"}
                        ]
                    },
                ),
            ],
            launch_step="Launching Rollers.Get sub-roller change IDs (1)",
        )
        + api.step_data(
            "Launching Rollers.Wait for Presubmit-Verified for roll CLs (0).Get CL details for CL 11111111",
            api.json.output(
                {
                    "current_revision": "abc",
                    "status": "NEW",
                    "labels": {
                        "Presubmit-Verified": {"approved": {}},
                    },
                }
            ),
        )
        + api.step_data(
            "Launching Rollers.Wait for Presubmit-Verified for roll CLs (0).Get CL details for CL 22222222",
            api.json.output(
                {
                    "current_revision": "abc",
                    "status": "NEW",
                    "labels": {
                        # Presubmit-Verified label doesn't have a vote yet!
                        "Presubmit-Verified": {},
                    },
                }
            ),
        )
        + api.step_data(
            "Launching Rollers.Wait for Presubmit-Verified for roll CLs (1).Get CL details for CL 22222222",
            api.json.output(
                {
                    "current_revision": "abc",
                    "status": "NEW",
                    "labels": {
                        "Presubmit-Verified": {"approved": {}},
                    },
                }
            ),
        )
    )

    yield (
        test("presubmit fails")
        + properties(builder("roller-1"))
        + simulate_subroller_launch(
            builds=[result_ci(name="roller-1", id=1000)],
        )
        + simulate_get_multi(
            builds=[
                result_ci(
                    name="roller-1",
                    id=1000,
                    output_props={
                        "gerrit_changes": [
                            {"host": "change_host_1", "change_id": "11111111"}
                        ]
                    },
                ),
            ],
            launch_step="Launching Rollers.Get sub-roller change IDs (0)",
        )
        + api.step_data(
            "Launching Rollers.Wait for Presubmit-Verified for roll CLs (0).Get CL details for CL 11111111",
            api.json.output(
                {
                    "current_revision": "abc",
                    "status": "NEW",
                    "labels": {
                        "Presubmit-Verified": {"rejected": {}},
                    },
                }
            ),
        )
    )

    yield (
        test("roll CL manually merged")
        + properties(builder("roller-1"))
        + simulate_subroller_launch(
            builds=[result_ci(name="roller-1", id=1000)],
        )
        + simulate_get_multi(
            builds=[
                result_ci(
                    name="roller-1",
                    id=1000,
                    output_props={
                        "gerrit_changes": [
                            {"host": "change_host_1", "change_id": "11111111"}
                        ]
                    },
                ),
            ],
            launch_step="Launching Rollers.Get sub-roller change IDs (0)",
        )
        + api.step_data(
            "Launching Rollers.Wait for Presubmit-Verified for roll CLs (0).Get CL details for CL 11111111",
            api.json.output(
                {
                    "current_revision": "abc",
                    "status": "MERGED",
                    "labels": {
                        "Presubmit-Verified": {"approved": {}},
                    },
                }
            ),
        )
    )

    yield (
        test("roll CL manually abandoned")
        + properties(builder("roller-1"))
        + simulate_subroller_launch(
            builds=[result_ci(name="roller-1", id=1000)],
        )
        + simulate_get_multi(
            builds=[
                result_ci(
                    name="roller-1",
                    id=1000,
                    output_props={
                        "gerrit_changes": [
                            {"host": "change_host_1", "change_id": "11111111"}
                        ]
                    },
                ),
            ],
            launch_step="Launching Rollers.Get sub-roller change IDs (0)",
        )
        + api.step_data(
            "Launching Rollers.Wait for Presubmit-Verified for roll CLs (0).Get CL details for CL 11111111",
            api.json.output(
                {
                    "current_revision": "abc",
                    "status": "ABANDONED",
                    "labels": {
                        "Presubmit-Verified": {"approved": {}},
                    },
                }
            ),
        )
    )

    yield (
        test("presubmit-verified label does not exist")
        + properties(builder("roller-1"))
        + simulate_subroller_launch(
            builds=[result_ci(name="roller-1", id=1000)],
        )
        + simulate_get_multi(
            builds=[
                result_ci(
                    name="roller-1",
                    id=1000,
                    output_props={
                        "gerrit_changes": [
                            {"host": "change_host_1", "change_id": "11111111"}
                        ]
                    },
                )
            ],
            launch_step="Launching Rollers.Get sub-roller change IDs (0)",
        )
        + api.step_data(
            "Launching Rollers.Wait for Presubmit-Verified for roll CLs (0).Get CL details for CL 11111111",
            api.json.output(
                {
                    "current_revision": "abc",
                    "status": "NEW",
                    "labels": {},
                }
            ),
        )
    )

    yield (
        test("get change id timeout", execution_timeout=10, status="FAILURE")
        + properties(builder("roller-1"))
        + simulate_subroller_launch(
            builds=[result_ci(name="roller-1", id=1000)],
        )
        + simulate_get_multi(
            builds=[
                result_ci(
                    name="roller-1",
                    id=1000,
                    output_props={},
                )
            ],
            launch_step="Launching Rollers.Get sub-roller change IDs (0)",
        )
    )

    yield (
        test("check for presubmit-verified timeout", execution_timeout=6)
        + properties(builder("roller-1"))
        + simulate_subroller_launch(
            builds=[result_ci(name="roller-1", id=1000)],
        )
        + simulate_get_multi(
            builds=[
                result_ci(
                    name="roller-1",
                    id=1000,
                    output_props={
                        "gerrit_changes": [
                            {"host": "change_host_1", "change_id": "11111111"}
                        ]
                    },
                )
            ],
            launch_step="Launching Rollers.Get sub-roller change IDs (0)",
        )
        + api.step_data(
            "Launching Rollers.Wait for Presubmit-Verified for roll CLs (0).Get CL details for CL 11111111",
            api.json.output(
                {
                    "current_revision": "abc",
                    "status": "NEW",
                    "labels": {
                        "Presubmit-Verified": {},
                    },
                }
            ),
        )
    )