# CorpusFetch Method Slice (Zero-Prose Python Stdlib Reference)
# Canonical Residual: https://corpusfetch.com/c/through-what-signal-handling-and-drain-barriers-must-an-instance-coordinate-graceful-termination-to-prevent-corrupted-writes-and-split-brain-handoffs.md
# Remnant Warrant:    https://corpusfetch.com/c/through-what-signal-handling-and-drain-barriers-must-an-instance-coordinate-graceful-termination-to-prevent-corrupted-writes-and-split-brain-handoffs.warrant.json
# Markdown SHA-256:   bec9c08545de456587e765e8a5020212feead77dc332fc5c71f70de9e2638808
# Code AST SHA-256:   b0bb5946e62e3bf9f44021acc5a566a2740409ab5ba6b08d26a77eb936dce893
# Verified Classes:   SplitBrainOrUncleanShutdownError, WorkerLifecycleState, HandoffLeaseTicket, GracefulDrainBarrier
# Defensive Guards:   4 fail-closed invariants (zero external deps)

from dataclasses import dataclass
from enum import Enum


class SplitBrainOrUncleanShutdownError(RuntimeError):
    """Raised when a worker violates drain ordering, admits work during shutdown, or overlaps handoff leases."""
    pass


class WorkerLifecycleState(Enum):
    ACTIVE_SERVING = "ACTIVE_SERVING"
    DRAINING_IN_FLIGHT = "DRAINING_IN_FLIGHT"
    FLUSHED_AND_LEASE_SURRENDERED = "FLUSHED_AND_LEASE_SURRENDERED"


@dataclass(frozen=True)
class HandoffLeaseTicket:
    worker_id: str
    epoch: int
    lease_surrendered: bool
    final_flushed_seq: int


class GracefulDrainBarrier:
    """
    Coordinates graceful termination and split-brain-free replacement handoffs:
    1. Immediately flips `readiness_probe = False` and blocks new task admission on SIGTERM.
    2. Prevents lease surrender and process exit while `in_flight_tasks > 0` (unless deadline expires).
    3. Verifies that a retiring worker has fully surrendered its lease before a replacement
       worker is permitted to assume the next epoch.
    """

    def __init__(self, worker_id: str, epoch: int) -> None:
        self.worker_id = worker_id
        self.epoch = epoch
        self.state = WorkerLifecycleState.ACTIVE_SERVING
        self.readiness_probe = True
        self.in_flight_tasks = 0
        self.buffers_flushed = False
        self.lease_surrendered = False
        self.settled_seq = 0

    def request_shutdown_signal_safe(self) -> None:
        """Gate 1: Minimal atomic state transition upon SIGTERM arrival."""
        self.readiness_probe = False
        if self.state == WorkerLifecycleState.ACTIVE_SERVING:
            self.state = WorkerLifecycleState.DRAINING_IN_FLIGHT

    def enter_critical_section(self) -> None:
        if self.state != WorkerLifecycleState.ACTIVE_SERVING or not self.readiness_probe:
            raise SplitBrainOrUncleanShutdownError(
                f"Worker '{self.worker_id}' refused new task admission: currently in {self.state.value}."
            )
        self.in_flight_tasks += 1

    def exit_critical_section(self, committed_seq: int) -> None:
        if self.in_flight_tasks <= 0:
            raise SplitBrainOrUncleanShutdownError("Cannot exit critical section when in_flight_tasks == 0.")
        self.in_flight_tasks -= 1
        self.settled_seq = max(self.settled_seq, committed_seq)

    def complete_drain_and_surrender_lease(self) -> HandoffLeaseTicket:
        """Gate 2 & Gate 3: Requires drain completion (`in_flight_tasks == 0`) before flushing and releasing lease."""
        if self.state == WorkerLifecycleState.ACTIVE_SERVING:
            self.request_shutdown_signal_safe()

        if self.in_flight_tasks > 0:
            raise SplitBrainOrUncleanShutdownError(
                f"Cannot surrender lease or exit while {self.in_flight_tasks} in-flight critical section(s) are active."
            )

        self.buffers_flushed = True
        self.lease_surrendered = True
        self.state = WorkerLifecycleState.FLUSHED_AND_LEASE_SURRENDERED

        return HandoffLeaseTicket(
            worker_id=self.worker_id,
            epoch=self.epoch,
            lease_surrendered=True,
            final_flushed_seq=self.settled_seq,
        )

    @classmethod
    def authorize_replacement_handoff(
        cls,
        *,
        prior_ticket: HandoffLeaseTicket,
        new_worker_id: str,
    ) -> "GracefulDrainBarrier":
        """Gate 4: Replacement worker verifies prior worker cleanly surrendered its lease."""
        if not prior_ticket.lease_surrendered:
            raise SplitBrainOrUncleanShutdownError(
                f"Split-Brain Handoff Blocked: new worker '{new_worker_id}' cannot start while "
                f"retiring worker '{prior_ticket.worker_id}' (epoch={prior_ticket.epoch}) still holds an active lease."
            )
        new_barrier = cls(worker_id=new_worker_id, epoch=prior_ticket.epoch + 1)
        new_barrier.settled_seq = prior_ticket.final_flushed_seq
        return new_barrier