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

from recipe_engine import recipe_api
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: float = 120) -> Generator[None, None, None]:
        current_time_sec: float = self.m.time.time()

        # Amount of time elapsed in the build.
        elapsed_time_sec: float = (
            current_time_sec - self.m.buildbucket.build.start_time.seconds
        )

        # Amount of time before the build times out.
        time_remaining_sec: float = (
            self.m.buildbucket.build.execution_timeout.seconds
            - elapsed_time_sec
        )

        # 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
        # GCS before the build times out.
        timeout_sec: float = max(time_remaining_sec - buffer_sec, 10.0)

        # Timing often looks weird in recipe unit tests. If the amount of time
        # remaining before the build times out is a negative number, ignore it
        # and proceed as if everything is internally consistent.
        if self._test_data.enabled and time_remaining_sec < 0:
            time_remaining_sec = timeout_sec = 10.0

        # Don't start steps if there's less than 10s remaining. A step timing
        # out looks like an infra failure which we don't want.
        if time_remaining_sec < 10.0:
            self.m.step.empty(
                'timed out',
                status=self.m.step.FAILURE,
                step_text=(
                    'The build has too little time remaining, not starting new '
                    f'steps. {time_remaining_sec}s remaining.'
                ),
            )

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

        try:
            with self.m.time.timeout(timeout_sec):
                yield

        finally:
            pass