# 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.
"""Provides a context manager to enforce timeouts on recipe steps.

This module helps ensure that critical recipe operations complete before the
overall build is terminated by Swarming due to exceeding its execution timeout.
It calculates the remaining time for the build and applies a shorter timeout
to the enclosed steps, allowing a buffer for cleanup or logging operations.

If the remaining build time is too short, it will prevent new steps from
starting and mark the build as a failure.
"""

from __future__ import annotations

import contextlib
from collections.abc import Generator

from recipe_engine import recipe_api

from RECIPE_MODULES.fuchsia.utils import nice_duration

DEFAULT_BUFFER_SECS = 120


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

    DEFAULT_BUFFER_SECS = DEFAULT_BUFFER_SECS

    @property
    def time_remaining_sec(self):
        """Amount of time before the build times out."""
        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.
        return (
            self.m.buildbucket.build.execution_timeout.seconds
            - elapsed_time_sec
        )

    def within_last(self, sec: float = DEFAULT_BUFFER_SECS) -> bool:
        time_remaining_sec = self.time_remaining_sec
        if self._test_data.enabled and time_remaining_sec <= 0:
            return False
        return time_remaining_sec < sec

    @contextlib.contextmanager
    def __call__(
        self,
        buffer_sec: float = DEFAULT_BUFFER_SECS,
    ) -> Generator[None, None, None]:
        time_remaining_sec = self.time_remaining_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