# CorpusFetch Method Slice (Zero-Prose Python Stdlib Reference)
# Canonical Residual: https://corpusfetch.com/c/across-what-monotonic-fencing-token-and-lease-heartbeat-gates-must-an-instance-execute-background-jobs-to-prevent-zombie-worker-split-brain-writes.md
# Remnant Warrant:    https://corpusfetch.com/c/across-what-monotonic-fencing-token-and-lease-heartbeat-gates-must-an-instance-execute-background-jobs-to-prevent-zombie-worker-split-brain-writes.warrant.json
# Markdown SHA-256:   16c1d3bbc881dc0ce06c23a76bebbafc3c53eecabad3b8ad2665520c06dcd73a
# Code AST SHA-256:   cdc62d515e38f61195cbae527cf47a0377595eede8141de334daa704cc1421a0
# Verified Classes:   ZombieWorkerFencedError, ClaimedJobLease, FencedJobStateMachine
# Defensive Guards:   3 fail-closed invariants (zero external deps)

import sqlite3
import time
from dataclasses import dataclass
from typing import Any


class ZombieWorkerFencedError(RuntimeError):
    '''
    Raised when a worker attempts to heartbeat, checkpoint, or complete a job
    whose `fencing_token` has already been superseded by a newer leaseholder.
    '''
    pass


@dataclass(frozen=True)
class ClaimedJobLease:
    job_id: str
    payload_json: str
    fencing_token: int
    worker_id: str
    locked_until_epoch: float


class FencedJobStateMachine:
    '''
    Implements a split-brain-proof background job queue in SQL:
    1. Atomically claims pending or expired-lease jobs while incrementing `fencing_token`.
    2. Verifies `fencing_token` on every heartbeat renewal and completion write.
    3. Rejects any write from a delayed "zombie" worker whose lease expired (`rowcount == 0`).
    '''

    def __init__(self, conn: sqlite3.Connection, *, lease_duration_sec: float = 30.0) -> None:
        if lease_duration_sec < 3.0:
            raise ValueError("lease_duration_sec must be at least 3.0 seconds")
        self.conn = conn
        self.lease_duration_sec = lease_duration_sec
        self._init_schema()

    def _init_schema(self) -> None:
        with self.conn:
            self.conn.execute(
                "CREATE TABLE IF NOT EXISTS fenced_jobs ("
                "  job_id TEXT PRIMARY KEY,"
                "  status TEXT NOT NULL CHECK(status IN ('PENDING', 'RUNNING', 'COMPLETED', 'FAILED')),"
                "  payload_json TEXT NOT NULL,"
                "  result_json TEXT,"
                "  fencing_token INTEGER NOT NULL DEFAULT 0,"
                "  worker_id TEXT,"
                "  locked_until_epoch REAL NOT NULL DEFAULT 0.0"
                ")"
            )

    def enqueue_job(self, job_id: str, payload_json: str) -> bool:
        '''Idempotently enqueues a new job with initial `fencing_token = 0`.'''
        with self.conn:
            cur = self.conn.execute(
                "INSERT OR IGNORE INTO fenced_jobs "
                "(job_id, status, payload_json, fencing_token, locked_until_epoch) "
                "VALUES (?, 'PENDING', ?, 0, 0.0)",
                (job_id, payload_json),
            )
            return cur.rowcount == 1

    def claim_next_job(
        self,
        worker_id: str,
        *,
        now_epoch: float | None = None,
    ) -> ClaimedJobLease | None:
        '''
        Gate 1: Atomically claims one PENDING job (or a RUNNING job whose lease has expired)
        and increments `fencing_token = fencing_token + 1` in a single atomic transaction.
        '''
        now = time.time() if now_epoch is None else now_epoch
        new_expiry = now + self.lease_duration_sec

        with self.conn:
            row = self.conn.execute(
                "UPDATE fenced_jobs "
                "SET status = 'RUNNING', "
                "    fencing_token = fencing_token + 1, "
                "    worker_id = ?, "
                "    locked_until_epoch = ? "
                "WHERE job_id = ("
                "  SELECT job_id FROM fenced_jobs "
                "  WHERE status = 'PENDING' "
                "     OR (status = 'RUNNING' AND locked_until_epoch < ?) "
                "  ORDER BY job_id ASC LIMIT 1"
                ") "
                "RETURNING job_id, payload_json, fencing_token, worker_id, locked_until_epoch",
                (worker_id, new_expiry, now),
            ).fetchone()

        if row is None:
            return None

        return ClaimedJobLease(
            job_id=row[0],
            payload_json=row[1],
            fencing_token=int(row[2]),
            worker_id=row[3],
            locked_until_epoch=float(row[4]),
        )

    def heartbeat_lease(
        self,
        lease: ClaimedJobLease,
        *,
        now_epoch: float | None = None,
    ) -> ClaimedJobLease:
        '''
        Gate 2: Extends `locked_until_epoch` ONLY if `fencing_token` still matches `lease.fencing_token`.
        Raises `ZombieWorkerFencedError` immediately if another worker stole the expired lease.
        '''
        now = time.time() if now_epoch is None else now_epoch
        new_expiry = now + self.lease_duration_sec

        with self.conn:
            cur = self.conn.execute(
                "UPDATE fenced_jobs "
                "SET locked_until_epoch = ? "
                "WHERE job_id = ? AND status = 'RUNNING' AND fencing_token = ?",
                (new_expiry, lease.job_id, lease.fencing_token),
            )
            if cur.rowcount == 0:
                raise ZombieWorkerFencedError(
                    f"Worker {lease.worker_id!r} lost lease on job {lease.job_id!r} "
                    f"(stale fencing_token={lease.fencing_token}); aborting immediately"
                )

        return ClaimedJobLease(
            job_id=lease.job_id,
            payload_json=lease.payload_json,
            fencing_token=lease.fencing_token,
            worker_id=lease.worker_id,
            locked_until_epoch=new_expiry,
        )

    def complete_job_fenced(
        self,
        lease: ClaimedJobLease,
        result_json: str,
    ) -> None:
        '''
        Gate 3: Commits final job completion strictly guarded by `fencing_token`.
        A paused zombie worker waking up after lease theft cannot overwrite the row.
        '''
        with self.conn:
            cur = self.conn.execute(
                "UPDATE fenced_jobs "
                "SET status = 'COMPLETED', result_json = ?, locked_until_epoch = 0.0 "
                "WHERE job_id = ? AND status = 'RUNNING' AND fencing_token = ?",
                (result_json, lease.job_id, lease.fencing_token),
            )
            if cur.rowcount == 0:
                raise ZombieWorkerFencedError(
                    f"Rejected zombie completion write for job {lease.job_id!r} "
                    f"with superseded fencing_token={lease.fencing_token}"
                )