#!/usr/bin/env python3
"""Original LockedIn Labs queue 101 lab. Python 3.11+, standard library only.

Run: python3 fde-queue-delivery-lab.py demo
     python3 fde-queue-delivery-lab.py test

Real SQLite writes; a deliberately small local queue model. Injected integer
ticks and omitted acknowledgements model interruptions, not actual process kills
or a cloud broker. All files live in temporary directories and are removed.
"""
import argparse
from contextlib import closing
import json
from pathlib import Path
import sqlite3
import tempfile
import unittest


class DeliveryLab:
    def __init__(self, path):
        self.path = path
        with closing(sqlite3.connect(path)) as db, db:
            db.executescript("""
                CREATE TABLE message (state TEXT, attempt INTEGER, until_tick INTEGER);
                INSERT INTO message VALUES ('ready', 0, 0);
                CREATE TABLE effects (tenant TEXT, operation TEXT, payload TEXT);
                CREATE TABLE receipts (tenant TEXT, operation TEXT, payload TEXT,
                                       PRIMARY KEY (tenant, operation));
            """)

    def receive(self, now):
        with closing(sqlite3.connect(self.path)) as db, db:
            db.execute("BEGIN IMMEDIATE")
            state, attempt, until = db.execute("SELECT * FROM message").fetchone()
            if state == "settled" or (state == "in flight" and now < until):
                raise ValueError("not_available")
            attempt += 1
            db.execute("UPDATE message SET state='in flight', attempt=?, until_tick=?", (attempt, now + 5))
            return attempt

    def acknowledge(self, receipt, now):
        with closing(sqlite3.connect(self.path)) as db, db:
            db.execute("BEGIN IMMEDIATE")
            state, attempt, until = db.execute("SELECT * FROM message").fetchone()
            if state != "in flight" or receipt != attempt or now >= until:
                raise ValueError("delivery_lost")
            db.execute("UPDATE message SET state='settled'")

    def apply(self, deduplicate, tenant="north", operation="J-12:draft:v1", payload="draft v1", fail=False):
        # This receipt and effect share ONE database transaction. This is not
        # sufficient for an email, payment or other external service's effect.
        with closing(sqlite3.connect(self.path)) as db, db:
            db.execute("BEGIN IMMEDIATE")
            if deduplicate:
                prior = db.execute("SELECT payload FROM receipts WHERE tenant=? AND operation=?", (tenant, operation)).fetchone()
                if prior:
                    if prior[0] != payload:
                        raise ValueError("operation_conflict")
                    return "reused existing effect"
                db.execute("INSERT INTO receipts VALUES (?, ?, ?)", (tenant, operation, payload))
            if fail:
                raise ValueError("interrupted_transaction")
            db.execute("INSERT INTO effects VALUES (?, ?, ?)", (tenant, operation, payload))
            return "created effect"

    def inspect(self):
        # A fresh connection observes committed state, not a cached Python value.
        with closing(sqlite3.connect(self.path)) as db:
            state, deliveries, _ = db.execute("SELECT * FROM message").fetchone()
            return {"queue": state, "deliveries": deliveries,
                    "effects": db.execute("SELECT count(*) FROM effects").fetchone()[0],
                    "receipts": db.execute("SELECT count(*) FROM receipts").fetchone()[0]}


def run_case(path, case):
    lab = DeliveryLab(path)
    trace = []

    def keep(event):
        trace.append({"event": event, **lab.inspect()})

    keep("1. Accepted J-12; no work done yet")
    first = lab.receive(now=0)
    keep("2. Worker A receives delivery 1; temporarily in flight")
    if case == "early-ack":
        lab.acknowledge(first, now=1)
        keep("3. A acknowledges BEFORE writing; model a stop here")
        try:
            lab.receive(now=5)
        except ValueError as error:
            keep(f"4. B cannot receive: {error}; the intended draft is missing")
        return trace

    safe = case != "duplicate-effect"
    lab.apply(deduplicate=safe)
    keep("3. A commits the local draft")
    if case == "normal":
        lab.acknowledge(first, now=1)
        keep("4. A acknowledges AFTER the write")
        return trace

    keep("4. Model a stop before acknowledgement; stored draft still exists")
    second = lab.receive(now=5)
    keep("5. Delivery 2 after expiry; same job and operation key")
    result = lab.apply(deduplicate=safe)
    keep(f"6. B {result}")
    lab.acknowledge(second, now=6)
    keep("7. B acknowledges its delivery")
    return trace


class DeliveryTests(unittest.TestCase):
    def setUp(self):
        self.temp = tempfile.TemporaryDirectory()
        self.addCleanup(self.temp.cleanup)
        self.path = Path(self.temp.name) / "lesson.sqlite"

    def test_normal(self):
        self.assertEqual(run_case(self.path, "normal")[-1]["effects"], 1)

    def test_repeat_without_receipt_duplicates_effect(self):
        last = run_case(self.path, "duplicate-effect")[-1]
        self.assertEqual((last["deliveries"], last["effects"]), (2, 2))

    def test_repeat_with_receipt_reuses_effect(self):
        last = run_case(self.path, "reuse-effect")[-1]
        self.assertEqual((last["queue"], last["deliveries"], last["effects"], last["receipts"]), ("settled", 2, 1, 1))

    def test_early_ack_loses_unperformed_work(self):
        last = run_case(self.path, "early-ack")[-1]
        self.assertEqual((last["queue"], last["effects"]), ("settled", 0))

    def test_expiry_and_old_delivery_receipt(self):
        lab = DeliveryLab(self.path)
        first = lab.receive(0)
        with self.assertRaisesRegex(ValueError, "not_available"):
            lab.receive(4)
        with self.assertRaisesRegex(ValueError, "delivery_lost"):
            lab.acknowledge(first, 5)
        second = lab.receive(5)
        with self.assertRaisesRegex(ValueError, "delivery_lost"):
            lab.acknowledge(first, 6)
        lab.acknowledge(second, 6)

    def test_conflicting_operation_is_not_a_replay_and_tenants_are_separate(self):
        lab = DeliveryLab(self.path)
        lab.apply(True)
        with self.assertRaisesRegex(ValueError, "operation_conflict"):
            lab.apply(True, payload="different draft")
        self.assertEqual(lab.inspect()["effects"], 1)
        lab.apply(True, tenant="south")
        self.assertEqual(lab.inspect()["effects"], 2)

    def test_receipt_rolls_back_with_effect(self):
        lab = DeliveryLab(self.path)
        with self.assertRaisesRegex(ValueError, "interrupted_transaction"):
            lab.apply(True, fail=True)
        self.assertEqual((lab.inspect()["receipts"], lab.inspect()["effects"]), (0, 0))
        self.assertEqual(lab.apply(True), "created effect")


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("command", choices=("demo", "test"))
    args = parser.parse_args()
    if args.command == "test":
        unittest.main(argv=[__file__])
    else:
        with tempfile.TemporaryDirectory(prefix="lockedin-queue-101-") as folder:
            for case in ("normal", "duplicate-effect", "reuse-effect", "early-ack"):
                print(json.dumps({"case": case, "trace": run_case(Path(folder) / f"{case}.sqlite", case)}, indent=2))
