# 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.
"""Ensure the given call completes before swarming kills the build."""

import contextlib
import dataclasses

from recipe_engine import recipe_api
from PB.go.chromium.org.luci.lucictx import sections as sections_pb2
from RECIPE_MODULES.fuchsia.utils import nice_duration


class TimeoutApi(recipe_api.RecipeApi):
    """Ensure the given call completes before swarming kills the build."""

    @contextlib.contextmanager
    def __call__(self, buffer_sec=120):
        current_time = self.m.time.time()

        # Amount of time elapsed in the run.
        elapsed_time = (
            current_time - self.m.buildbucket.build.start_time.seconds
        )

        # Amount of time before build times out.
        time_remaining = (
            self.m.buildbucket.build.execution_timeout.seconds - elapsed_time
        )

        # Give a buffer before build times out and kill this step then. This
        # should give enough time to read any logfiles and maybe upload to
        # logdog/GCS before the build times out.
        timeout = time_remaining - buffer_sec

        soft_deadline = current_time + timeout

        with self.m.step.nest(f'timeout {nice_duration(timeout)}'):
            pass

        deadline = sections_pb2.Deadline()
        deadline.soft_deadline = soft_deadline
        deadline.grace_period = 30.0

        try:
            with self.m.context(deadline=deadline):
                yield

        finally:
            pass