#!/usr/bin/env python3
"""Original LockedIn Labs synthetic SQLite worker-lease exercise, Python 3.11+.
Run: python3 fde-worker-recovery-lab.py test
     python3 fde-worker-recovery-lab.py demo
     python3 fde-worker-recovery-lab.py crash-demo  # macOS/Linux only
     python3 fde-worker-recovery-lab.py race-demo
No model, queue service, network, credentials or external effect is used.
Injected clock ticks make the failure cases deterministic, not a load benchmark.
"""
import argparse
from dataclasses import dataclass
import hashlib
import json
import multiprocessing
import os
from pathlib import Path
import selectors
import signal
import socket
import sqlite3
import subprocess
import sys
import tempfile
import threading
import unittest
from unittest.mock import patch


def digest(value):
    return hashlib.sha256(value.encode()).hexdigest()


@dataclass(frozen=True)
class Lease:
    tenant: str
    request: str
    worker: str
    fence: int


class JobQueue:
    """Local teaching reference. Host adapters must supply trusted identity/time.
    Fencing protects this completion record, not an external destination.
    """
    def __init__(self, path):
        self.db = sqlite3.connect(path, timeout=1)
        self.db.row_factory = sqlite3.Row
        self.db.execute("""CREATE TABLE IF NOT EXISTS jobs (
            tenant TEXT, request TEXT, payload_digest TEXT, checkpoint_version INTEGER,
            state TEXT, worker TEXT, fence INTEGER, lease_until INTEGER,
            attempts INTEGER, attempt_limit INTEGER, result_digest TEXT,
            PRIMARY KEY (tenant, request))""")

    def close(self):
        self.db.close()

    def inspect(self, tenant, request):
        row = self.db.execute("SELECT * FROM jobs WHERE tenant=? AND request=?", (tenant, request)).fetchone()
        if row is None:
            raise PermissionError("job_unavailable")
        return dict(row)

    def admit(self, tenant, request, payload, checkpoint_version=1):
        if not tenant or not request or len(request) > 80 or checkpoint_version not in (1, 2):
            raise ValueError("invalid_admission")
        with self.db:
            self.db.execute("BEGIN IMMEDIATE")
            prior = self.db.execute("SELECT * FROM jobs WHERE tenant=? AND request=?", (tenant, request)).fetchone()
            if prior:
                if prior["payload_digest"] != digest(payload) or prior["checkpoint_version"] != checkpoint_version:
                    raise ValueError("request_conflict")
                return dict(prior)
            active = self.db.execute("SELECT count(*) FROM jobs WHERE tenant=? AND state IN ('queued','running')", (tenant,)).fetchone()[0]
            if active >= 2:
                raise ValueError("admission_full")
            self.db.execute("INSERT INTO jobs VALUES (?, ?, ?, ?, 'queued', NULL, 0, NULL, 0, 2, NULL)",
                            (tenant, request, digest(payload), checkpoint_version))
        return self.inspect(tenant, request)

    def claim(self, tenant, request, worker, now, lease_ticks=5, supported_versions=(1,)):
        if not worker or now < 0 or not 1 <= lease_ticks <= 10:
            raise ValueError("invalid_lease")
        with self.db:
            self.db.execute("BEGIN IMMEDIATE")
            job = self.inspect(tenant, request)
            if job["state"] not in ("queued", "running"):
                raise ValueError("not_claimable")
            if job["checkpoint_version"] not in supported_versions:
                raise ValueError("checkpoint_incompatible")
            if job["state"] == "running" and now < job["lease_until"]:
                raise ValueError("lease_busy")
            if job["attempts"] >= job["attempt_limit"]:
                raise ValueError("attempt_budget_exhausted")
            fence = job["fence"] + 1
            self.db.execute("UPDATE jobs SET state='running', worker=?, fence=?, lease_until=?, attempts=attempts+1 WHERE tenant=? AND request=?",
                            (worker, fence, now + lease_ticks, tenant, request))
        return Lease(tenant, request, worker, fence)

    def finish(self, lease, result, now):
        with self.db:
            self.db.execute("BEGIN IMMEDIATE")
            job = self.inspect(lease.tenant, lease.request)
            if (job["worker"], job["fence"]) != (lease.worker, lease.fence):
                raise ValueError("lease_lost")
            if job["state"] == "finished":
                if job["result_digest"] != digest(result):
                    raise ValueError("result_conflict")
                return job  # Receipt replay only: no work is repeated.
            if job["state"] != "running" or now >= job["lease_until"]:
                raise ValueError("lease_lost")
            self.db.execute("UPDATE jobs SET state='finished', result_digest=? WHERE tenant=? AND request=?",
                            (digest(result), lease.tenant, lease.request))
        return self.inspect(lease.tenant, lease.request)

    def cancel(self, tenant, request):
        with self.db:
            self.db.execute("BEGIN IMMEDIATE")
            job = self.inspect(tenant, request)
            if job["state"] == "finished":
                raise ValueError("already_finished")
            if job["state"] != "cancelled":
                self.db.execute("UPDATE jobs SET state='cancelled', worker=NULL, fence=fence+1, lease_until=NULL WHERE tenant=? AND request=?",
                                (tenant, request))
        return self.inspect(tenant, request)


class QueueTests(unittest.TestCase):
    def setUp(self):
        self.folder = tempfile.TemporaryDirectory()
        self.path = str(Path(self.folder.name) / "synthetic.sqlite")
        self.queue = JobQueue(self.path)

    def tearDown(self):
        self.queue.close()
        self.folder.cleanup()

    def test_admission_is_not_completion_and_replay_is_bound(self):
        first = self.queue.admit("north", "R1", "synthetic lookup")
        self.assertEqual(first["state"], "queued")
        self.assertEqual(first, self.queue.admit("north", "R1", "synthetic lookup"))
        with self.assertRaisesRegex(ValueError, "request_conflict"):
            self.queue.admit("north", "R1", "changed work")
        with self.assertRaisesRegex(PermissionError, "job_unavailable"):
            self.queue.inspect("south", "R1")

    def test_admission_capacity_is_tenant_scoped(self):
        for request in ("R1", "R2"):
            self.queue.admit("north", request, "synthetic")
        with self.assertRaisesRegex(ValueError, "admission_full"):
            self.queue.admit("north", "R3", "synthetic")
        self.assertEqual(self.queue.admit("south", "R3", "synthetic")["state"], "queued")

    def test_abandoned_lease_survives_connection_restart_and_old_worker_is_fenced(self):
        self.queue.admit("north", "R1", "synthetic")
        old = self.queue.claim("north", "R1", "worker-a", now=0)
        self.queue.close()  # Simulate a lost worker; no completion receipt exists.
        self.queue = JobQueue(self.path)
        current = self.queue.claim("north", "R1", "worker-b", now=5)
        self.assertEqual((old.fence, current.fence), (1, 2))
        with self.assertRaisesRegex(ValueError, "lease_lost"):
            self.queue.finish(old, "old result", now=6)
        self.assertEqual(self.queue.finish(current, "current result", now=6)["state"], "finished")

    def test_live_lease_cannot_be_stolen_and_expired_worker_cannot_finish(self):
        self.queue.admit("north", "R1", "synthetic")
        lease = self.queue.claim("north", "R1", "worker-a", now=0)
        with self.assertRaisesRegex(ValueError, "lease_busy"):
            self.queue.claim("north", "R1", "worker-b", now=4)
        with self.assertRaisesRegex(ValueError, "lease_lost"):
            self.queue.finish(lease, "late result", now=5)

    def test_cancellation_revokes_recording_and_replay_does_not_resurrect(self):
        self.queue.admit("north", "R1", "synthetic")
        lease = self.queue.claim("north", "R1", "worker-a", now=0)
        cancelled = self.queue.cancel("north", "R1")
        self.assertEqual(cancelled, self.queue.cancel("north", "R1"))
        self.assertEqual(self.queue.admit("north", "R1", "synthetic")["state"], "cancelled")
        with self.assertRaisesRegex(ValueError, "lease_lost"):
            self.queue.finish(lease, "late result", now=1)

    def test_checkpoint_compatibility_is_checked_before_attempt_is_consumed(self):
        self.queue.admit("north", "R1", "synthetic", checkpoint_version=2)
        with self.assertRaisesRegex(ValueError, "checkpoint_incompatible"):
            self.queue.claim("north", "R1", "old-worker", now=0)
        self.assertEqual(self.queue.inspect("north", "R1")["attempts"], 0)
        self.queue.claim("north", "R1", "new-worker", now=0, supported_versions=(2,))
        self.assertEqual(self.queue.inspect("north", "R1")["attempts"], 1)

    def test_attempt_budget_survives_restart(self):
        self.queue.admit("north", "R1", "synthetic")
        self.queue.claim("north", "R1", "worker-a", now=0)
        self.queue.claim("north", "R1", "worker-b", now=5)
        self.queue.close()
        self.queue = JobQueue(self.path)
        with self.assertRaisesRegex(ValueError, "attempt_budget_exhausted"):
            self.queue.claim("north", "R1", "worker-c", now=10)
        self.assertEqual(self.queue.inspect("north", "R1")["attempts"], 2)

    def test_completion_receipt_replay_cannot_change_result(self):
        self.queue.admit("north", "R1", "synthetic")
        lease = self.queue.claim("north", "R1", "worker-a", now=0)
        receipt = self.queue.finish(lease, "synthetic result", now=1)
        self.assertEqual(receipt, self.queue.finish(lease, "synthetic result", now=20))
        with self.assertRaisesRegex(ValueError, "result_conflict"):
            self.queue.finish(lease, "different result", now=20)


def hold_worker(path):
    """Child-only entry: commit a lease, report the boundary, then stay alive."""
    # Never open an existing file, including a caller's real database.
    with Path(path).open("xb"):
        pass
    queue = JobQueue(path)
    try:
        queue.admit("north", "R1", "synthetic lookup")
        queue.claim("north", "R1", "worker-a", now=0)
        print(json.dumps({"event": "lease_committed", "job": queue.inspect("north", "R1")}), flush=True)
        threading.Event().wait()  # No transaction or external action is in flight.
    finally:
        queue.close()  # SIGKILL intentionally does not execute this cleanup.


def crash_rehearsal():
    """Kill only the child created here, after its named committed boundary."""
    if os.name != "posix":
        raise RuntimeError("crash_demo_requires_macos_or_linux")
    with tempfile.TemporaryDirectory(prefix="field-worker-crash-") as folder:
        path = str(Path(folder) / "synthetic.sqlite")
        child = subprocess.Popen(
            [sys.executable, str(Path(__file__).resolve()), "_hold-worker", "--db", path],
            stdin=subprocess.DEVNULL, stdout=subprocess.PIPE, stderr=subprocess.PIPE,
            text=True, shell=False,
        )
        try:
            with selectors.DefaultSelector() as ready:
                ready.register(child.stdout, selectors.EVENT_READ)
                if not ready.select(timeout=10):
                    raise RuntimeError("worker_boundary_not_observed")
                boundary = json.loads(child.stdout.readline())
            if boundary.get("event") != "lease_committed":
                raise RuntimeError("unexpected_worker_boundary")
            child.kill()
            exit_code = child.wait(timeout=5)
            if exit_code != -signal.SIGKILL:
                raise RuntimeError("sigkill_not_observed")
            reports = {"boundary": boundary, "worker_exit": exit_code}
            queue = JobQueue(path)  # A different process now reads the committed lease.
            try:
                reports["after_kill"] = queue.inspect("north", "R1")
                reports["integrity_check"] = queue.db.execute("PRAGMA integrity_check").fetchone()[0]
                try:
                    queue.claim("north", "R1", "worker-b", now=4)
                except ValueError as error:
                    reports["before_expiry"] = str(error)
                reports["attempts_before_reclaim"] = queue.inspect("north", "R1")["attempts"]
                current = queue.claim("north", "R1", "worker-b", now=5)
                reports["reclaimed"] = queue.inspect("north", "R1")
                # Reconstruct a stale submission; the killed process cannot send it.
                stale = Lease("north", "R1", "worker-a", boundary["job"]["fence"])
                try:
                    queue.finish(stale, "stale output", now=6)
                except ValueError as error:
                    reports["reconstructed_stale_submission"] = str(error)
                receipt = queue.finish(current, "synthetic output", now=6)
                reports["finished"] = receipt
                reports["receipt_replay_equal"] = receipt == queue.finish(current, "synthetic output", now=20)
                reports["environment"] = {"python": sys.version.split()[0], "sqlite": sqlite3.sqlite_version}
                reports["limits"] = "SIGKILL after a committed lease, while idle. Injected ticks, not measured expiry. No mid-transaction kill, power loss, remote queue, external effect, model or authentication is tested."
                return reports
            finally:
                queue.close()
        finally:
            if child.poll() is None:
                child.kill()
            child.communicate(timeout=5)  # Reap the owned child even on a failed assertion.


def competing_claim(path, worker, ready, results):
    """Owned child: open its own connection before the shared starting gate."""
    with patch.object(socket.socket, "connect", side_effect=AssertionError("network_forbidden")):
        queue = JobQueue(path)
        try:
            ready.wait(timeout=10)
            try:
                lease = queue.claim("north", "R1", worker, now=0)
                result = {"code": "claimed", "fence": lease.fence}
            except ValueError as error:
                result = {"code": str(error)}
            results.put({"worker": worker, "pid": os.getpid(), **result})
        finally:
            queue.close()


def race_rehearsal():
    """Release two fresh local workers against one already admitted job."""
    with tempfile.TemporaryDirectory(prefix="field-worker-race-") as folder:
        path = str(Path(folder) / "synthetic.sqlite")
        queue = JobQueue(path)
        try:
            admitted = queue.admit("north", "R1", "synthetic lookup")
        finally:
            queue.close()
        context = multiprocessing.get_context("spawn")
        ready = context.Barrier(3)
        results = context.Queue()
        children = [context.Process(target=competing_claim, args=(path, worker, ready, results))
                    for worker in ("worker-a", "worker-b")]
        started = []
        try:
            for child in children:
                child.start()
                started.append(child)
            ready.wait(timeout=10)  # Both children have opened separate connections.
            observations = [results.get(timeout=10) for _ in children]
            for child in started:
                child.join(timeout=5)
                if child.is_alive() or child.exitcode != 0:
                    raise RuntimeError("claim_worker_did_not_finish")
            queue = JobQueue(path)
            try:
                final = queue.inspect("north", "R1")
                return {
                    "admitted": admitted,
                    "workers": sorted(observations, key=lambda item: item["worker"]),
                    "after_claims": final,
                    "integrity_check": queue.db.execute("PRAGMA integrity_check").fetchone()[0],
                    "environment": {"python": sys.version.split()[0], "sqlite": sqlite3.sqlite_version},
                    "limits": "Two owned local processes released from one barrier; execution overlap and winner ordering are not guaranteed. One local SQLite file, injected tick 0. No throughput, external effect, multi-host queue, clock agreement or authentication is tested.",
                }
            finally:
                queue.close()
        finally:
            for child in started:
                if child.is_alive():
                    child.kill()
                child.join(timeout=5)
            results.close()
            results.join_thread()


class ConcurrentClaimTests(unittest.TestCase):
    def test_two_processes_share_one_lease_without_consuming_two_attempts(self):
        report = race_rehearsal()
        self.assertEqual(report["admitted"]["state"], "queued")
        workers = report["workers"]
        self.assertEqual(len({worker["pid"] for worker in workers}), 2)
        self.assertNotIn(os.getpid(), [worker["pid"] for worker in workers])
        self.assertCountEqual([worker["code"] for worker in workers], ["claimed", "lease_busy"])
        winner = next(worker for worker in workers if worker["code"] == "claimed")
        final = report["after_claims"]
        self.assertEqual((winner["fence"], final["fence"], final["attempts"]), (1, 1, 1))
        self.assertEqual(final["worker"], winner["worker"])
        self.assertEqual((final["state"], final["lease_until"]), ("running", 5))
        self.assertIsNone(final["result_digest"])
        self.assertEqual(report["integrity_check"], "ok")


@unittest.skipUnless(os.name == "posix", "SIGKILL rehearsal requires macOS/Linux")
class ProcessRecoveryTests(unittest.TestCase):
    def test_actual_killed_process_leaves_a_lease_that_can_be_recovered(self):
        report = crash_rehearsal()
        self.assertEqual(report["worker_exit"], -signal.SIGKILL)
        self.assertEqual(report["boundary"]["job"], report["after_kill"])
        self.assertEqual(report["after_kill"]["state"], "running")
        self.assertIsNone(report["after_kill"]["result_digest"])
        self.assertEqual(report["integrity_check"], "ok")
        self.assertEqual(report["before_expiry"], "lease_busy")
        self.assertEqual(report["attempts_before_reclaim"], 1)
        self.assertEqual((report["reclaimed"]["fence"], report["reclaimed"]["attempts"]), (2, 2))
        self.assertEqual(report["reconstructed_stale_submission"], "lease_lost")
        self.assertEqual(report["finished"]["state"], "finished")
        self.assertTrue(report["receipt_replay_equal"])

    def test_child_entry_preserves_an_existing_file(self):
        with tempfile.TemporaryDirectory() as folder:
            path = Path(folder) / "keep.sqlite"
            path.write_bytes(b"preserve this file")
            for command, error in (("_hold-worker", "new_file_required"),
                                   ("crash-demo", "db_argument_is_internal_only"),
                                   ("race-demo", "db_argument_is_internal_only")):
                with self.subTest(command=command):
                    result = subprocess.run(
                        [sys.executable, str(Path(__file__).resolve()), command, "--db", str(path)],
                        capture_output=True, text=True, timeout=5, shell=False,
                    )
                    self.assertEqual(result.returncode, 2)
                    self.assertIn(error, result.stderr)
                    self.assertEqual(path.read_bytes(), b"preserve this file")


def demo():
    with tempfile.TemporaryDirectory(prefix="field-worker-") as folder:
        path = str(Path(folder) / "synthetic.sqlite")
        queue = JobQueue(path)
        reports = {"admitted": queue.admit("north", "R1", "synthetic lookup")}
        old = queue.claim("north", "R1", "worker-a", now=0)
        reports["abandoned"] = queue.inspect("north", "R1")
        queue.close()
        queue = JobQueue(path)
        try:
            current = queue.claim("north", "R1", "worker-b", now=5)
            try:
                queue.finish(old, "stale output", now=6)
            except ValueError as error:
                reports["old_worker_completion"] = str(error)
            reports["reclaimed"] = queue.inspect("north", "R1")
            reports["finished"] = queue.finish(current, "synthetic output", now=6)
            reports["limits"] = "Local simulated lease loss; no external action, process kill, distributed queue, authentication or customer acceptance is demonstrated."
            print(json.dumps(reports, indent=2))
        finally:
            queue.close()


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("command", choices=["test", "demo", "crash-demo", "race-demo", "_hold-worker"])
    parser.add_argument("--db", help=argparse.SUPPRESS)
    args = parser.parse_args()
    if args.db is not None and args.command != "_hold-worker":
        parser.error("db_argument_is_internal_only")
    with patch.object(socket.socket, "connect", side_effect=AssertionError("network_forbidden")):
        if args.command == "test":
            suite = unittest.TestSuite(unittest.defaultTestLoader.loadTestsFromTestCase(case)
                                       for case in (QueueTests, ProcessRecoveryTests, ConcurrentClaimTests))
            result = unittest.TextTestRunner(verbosity=2).run(suite)
            raise SystemExit(0 if result.wasSuccessful() else 1)
        elif args.command == "race-demo":
            print(json.dumps(race_rehearsal(), indent=2))
        elif args.command == "crash-demo":
            if os.name != "posix":
                parser.error("crash_demo_requires_macos_or_linux")
            print(json.dumps(crash_rehearsal(), indent=2))
        elif args.command == "_hold-worker":
            if not args.db:
                parser.error("new_disposable_db_path_required")
            try:
                hold_worker(args.db)
            except FileExistsError:
                parser.error("new_file_required")
        else:
            demo()
