# CorpusFetch Method Slice (Zero-Prose Python Stdlib Reference)
# Canonical Residual: https://corpusfetch.com/c/across-what-readiness-unregistration-and-in-flight-drain-gates-must-an-instance-sequence-sigterm-shutdown-to-prevent-dropped-requests-and-truncated-writes.md
# Remnant Warrant:    https://corpusfetch.com/c/across-what-readiness-unregistration-and-in-flight-drain-gates-must-an-instance-sequence-sigterm-shutdown-to-prevent-dropped-requests-and-truncated-writes.warrant.json
# Markdown SHA-256:   7bc370b2c24b9a4bbcace1e8a896ffd5b1e4eb711e7b28adb16c7079e53c4885
# Code AST SHA-256:   78c82a10213d841e2ffb4861183f58f481d2373dfadae3b3b541da645b993f8d
# Verified Classes:   ServiceShuttingDownError, GracefulDrainController
# Defensive Guards:   1 fail-closed invariants (zero external deps)

import asyncio
import signal
import time
from contextlib import asynccontextmanager
from typing import AsyncIterator, Awaitable, Callable


class ServiceShuttingDownError(RuntimeError):
    '''Raised when new work is rejected after the drain barrier has closed.'''
    pass


class GracefulDrainController:
    '''
    Coordinates a four-phase graceful shutdown on SIGTERM / SIGINT:
    1. Flips `is_ready = False` immediately so `/readyz` returns HTTP 503.
    2. Waits an optional load-balancer deregistration window (`lb_propagation_sec`)
       while still admitting edge requests in flight.
    3. Seals the admission gate (`is_accepting = False`) and awaits all tracked
       in-flight tasks up to `drain_timeout_sec`.
    4. Executes outbound resource cleanup hooks (DB pools, caches, telemetry flush)
       strictly in reverse registration order after in-flight work reaches zero.
    '''

    def __init__(
        self,
        *,
        lb_propagation_sec: float = 2.0,
        drain_timeout_sec: float = 20.0,
    ) -> None:
        self.lb_propagation_sec = max(0.0, lb_propagation_sec)
        self.drain_timeout_sec = max(1.0, drain_timeout_sec)
        self.is_ready: bool = True
        self.is_accepting: bool = True
        self._in_flight_count: int = 0
        self._shutdown_triggered = asyncio.Event()
        self._zero_in_flight = asyncio.Event()
        self._zero_in_flight.set()
        self._cleanup_callbacks: list[Callable[[], Awaitable[None]]] = []

    def install_signal_handlers(self, loop: asyncio.AbstractEventLoop) -> None:
        '''
        Registers async-safe signal handlers via loop.add_signal_handler
        (never doing I/O or locking inside a raw C signal frame).
        '''
        for sig in (signal.SIGTERM, signal.SIGINT):
            try:
                loop.add_signal_handler(sig, self.initiate_shutdown)
            except NotImplementedError:
                signal.signal(
                    sig,
                    lambda _s, _f: loop.call_soon_threadsafe(self.initiate_shutdown),
                )

    def initiate_shutdown(self) -> None:
        '''Phase 1: Immediately mark readiness false and wake the drain coordinator.'''
        self.is_ready = False
        self._shutdown_triggered.set()

    def register_cleanup_hook(self, async_close_fn: Callable[[], Awaitable[None]]) -> None:
        '''Registers an outbound pool/client teardown callback run only in Phase 4.'''
        self._cleanup_callbacks.append(async_close_fn)

    @asynccontextmanager
    async def track_in_flight_work(self) -> AsyncIterator[None]:
        '''
        Wraps every request handler or queue job. Rejects new work once Phase 3
        starts and signals `_zero_in_flight` the moment the last active task exits.
        '''
        if not self.is_accepting:
            raise ServiceShuttingDownError("Service is draining and no longer accepting new work")

        self._in_flight_count += 1
        self._zero_in_flight.clear()
        try:
            yield
        finally:
            self._in_flight_count -= 1
            if self._in_flight_count <= 0:
                self._in_flight_count = 0
                self._zero_in_flight.set()

    async def wait_for_shutdown_and_drain(self) -> dict[str, float | int]:
        '''
        Awaits the shutdown signal, then executes Phases 2, 3, and 4 in strict sequence.
        '''
        await self._shutdown_triggered.wait()
        t_start = time.monotonic()

        # Phase 2: Keep accepting in-flight load-balancer traffic while /readyz returns 503
        if self.lb_propagation_sec > 0:
            await asyncio.sleep(self.lb_propagation_sec)

        # Phase 3: Seal admission gate and wait for active tasks to finish
        self.is_accepting = False
        timed_out_tasks = 0
        try:
            await asyncio.wait_for(
                self._zero_in_flight.wait(),
                timeout=self.drain_timeout_sec,
            )
        except asyncio.TimeoutError:
            timed_out_tasks = self._in_flight_count

        # Phase 4: Close outbound DB pools, Redis clients, and telemetry in LIFO order
        for close_fn in reversed(self._cleanup_callbacks):
            try:
                await asyncio.wait_for(close_fn(), timeout=3.0)
            except Exception:
                pass

        return {
            "drain_elapsed_sec": round(time.monotonic() - t_start, 3),
            "abandoned_in_flight_tasks": timed_out_tasks,
        }