#!/usr/bin/env python3
"""Original LockedIn Labs teaching lab. Fictional data; SQLite FTS5; no network.

Python 3.11+ with SQLite FTS5. All changes occur in a disposable memory database.
The web visual replays `export`; this file actually executes the SQL.
"""
import argparse
import json
import sqlite3
import unittest


OLD_TEXT = "Handoff procedure: use the old checklist and record the desk owner."
NEW_TEXT = "Handoff procedure: use the revised checklist and obtain owner review."
STAGES = [
    ("current", "Start with a current copy", "The source and its searchable copy both describe revision 1."),
    ("corrected", "Correct the source", "Revision 2 is saved. The indexing job has not run yet."),
    ("refreshed", "Refresh the index", "The worker replaces the old searchable copy with revision 2."),
    ("late-job", "Replay the old job", "A delayed revision 1 job arrives after revision 2. The worker rejects it."),
    ("deleted", "Delete the source", "Revision 3 is a deletion marker. The searchable copy still contains revision 2."),
    ("purged", "Purge the search copy", "The deletion job removes the indexed row. The source deletion marker remains."),
    ("replayed", "Try to resurrect it", "A revision 2 job is retried after deletion. It cannot recreate the search copy."),
]


def database():
    connection = sqlite3.connect(":memory:")
    connection.row_factory = sqlite3.Row
    connection.execute("CREATE TABLE sources (id TEXT PRIMARY KEY, revision INTEGER, state TEXT, body TEXT)")
    try:
        connection.execute("CREATE VIRTUAL TABLE search USING fts5(source_id UNINDEXED, revision UNINDEXED, body)")
    except sqlite3.OperationalError as error:
        connection.close()
        raise RuntimeError("This Python installation needs SQLite with FTS5 enabled. No extension is downloaded by this lab.") from error
    with connection:
        connection.execute("INSERT INTO sources VALUES (?, ?, ?, ?)", ("P-17", 1, "active", OLD_TEXT))
    return connection


def apply_job(connection, revision, operation="upsert", source_id="P-17"):
    """Jobs name a version; the worker reads content from the authoritative row.

    This local transaction is not a distributed transaction or queue service.
    """
    with connection:
        source = connection.execute("SELECT * FROM sources WHERE id = ?", (source_id,)).fetchone()
        if not source:
            return "ignored: source missing"
        if source["revision"] != revision:
            return "ignored: revision does not match source"
        if operation == "delete" and source["state"] == "deleted":
            count = connection.execute("SELECT count(*) FROM search WHERE source_id = ?", (source_id,)).fetchone()[0]
            connection.execute("DELETE FROM search WHERE source_id = ?", (source_id,))
            return "purged" if count else "unchanged: already purged"
        if operation != "upsert" or source["state"] != "active":
            return "ignored: operation does not match source state"
        copies = connection.execute("SELECT revision, body FROM search WHERE source_id = ?", (source_id,)).fetchall()
        if len(copies) == 1 and int(copies[0]["revision"]) == revision and copies[0]["body"] == source["body"]:
            return "unchanged: current copy already present"
        connection.execute("DELETE FROM search WHERE source_id = ?", (source_id,))
        connection.execute("INSERT INTO search (source_id, revision, body) VALUES (?, ?, ?)", (source_id, revision, source["body"]))
        return "indexed: revision " + str(revision)


def retrieve(connection, query):
    # Treat input as one literal FTS phrase, not arbitrary FTS syntax or SQL.
    phrase = '"' + query.replace('"', '""') + '"'
    matches = connection.execute("SELECT source_id, revision FROM search WHERE search MATCH ? ORDER BY rowid", (phrase,)).fetchall() if query.strip() else []
    decisions, context = [], []
    for match in matches:
        source = connection.execute("SELECT * FROM sources WHERE id = ?", (match["source_id"],)).fetchone()
        if not source:
            reason = "missing"
        elif source["state"] != "active":
            reason = "deleted"
        elif source["revision"] != int(match["revision"]):
            reason = "stale"
        else:
            reason = "current"
            # Only current source text crosses into the context packet.
            context.append({"id": source["id"], "revision": source["revision"], "body": source["body"]})
        decisions.append({"id": match["source_id"], "indexedRevision": int(match["revision"]), "reason": reason})
    return {"matches": decisions, "context": context}


def snapshot(connection, stage, worker, query):
    identifier, label, event = STAGES[stage]
    source = dict(connection.execute("SELECT * FROM sources WHERE id = 'P-17'").fetchone())
    indexed = [dict(row) for row in connection.execute("SELECT revision, body FROM search WHERE source_id = 'P-17' ORDER BY rowid")]
    for row in indexed:
        row["revision"] = int(row["revision"])
    return {"id": identifier, "label": label, "event": event, "worker": worker,
            "source": source, "indexed": indexed, **retrieve(connection, query)}


def timeline(query="handoff"):
    connection = database()
    try:
        result = [snapshot(connection, 0, apply_job(connection, 1), query)]
        with connection:
            connection.execute("UPDATE sources SET revision = 2, body = ? WHERE id = 'P-17'", (NEW_TEXT,))
        result.append(snapshot(connection, 1, "waiting for revision 2 job", query))
        result.append(snapshot(connection, 2, apply_job(connection, 2), query))
        result.append(snapshot(connection, 3, apply_job(connection, 1), query))
        with connection:
            connection.execute("UPDATE sources SET revision = 3, state = 'deleted', body = '' WHERE id = 'P-17'")
        result.append(snapshot(connection, 4, "waiting for deletion job", query))
        result.append(snapshot(connection, 5, apply_job(connection, 3, "delete"), query))
        result.append(snapshot(connection, 6, apply_job(connection, 2), query))
        return {"engine": "SQLite FTS5 · recorded local execution", "query": query, "stages": result}
    finally:
        connection.close()


class ProofTests(unittest.TestCase):
    def test_stale_search_hit_is_not_context(self):
        corrected = timeline()["stages"][1]
        self.assertEqual(corrected["matches"], [{"id": "P-17", "indexedRevision": 1, "reason": "stale"}])
        self.assertEqual(corrected["context"], [])

    def test_refresh_replaces_searchable_terms(self):
        old_query = timeline("old")["stages"]
        self.assertEqual(len(old_query[0]["matches"]), 1)
        self.assertEqual(old_query[2]["matches"], [])
        self.assertEqual(timeline("revised")["stages"][2]["context"][0]["revision"], 2)

    def test_late_job_cannot_roll_back_current_copy(self):
        stages = timeline()["stages"]
        self.assertEqual(stages[3]["indexed"], stages[2]["indexed"])
        self.assertTrue(stages[3]["worker"].startswith("ignored"))

    def test_delete_blocks_context_before_purge_and_replay_cannot_resurrect(self):
        stages = timeline()["stages"]
        self.assertEqual(stages[4]["matches"][0]["reason"], "deleted")
        for stage in stages[4:]:
            self.assertEqual(stage["context"], [])
        self.assertEqual(stages[6]["indexed"], [])
        self.assertEqual(stages[6]["source"]["revision"], 3)

    def test_duplicate_job_does_not_duplicate_rows(self):
        connection = database()
        try:
            apply_job(connection, 1)
            self.assertEqual(apply_job(connection, 1), "unchanged: current copy already present")
            self.assertEqual(connection.execute("SELECT count(*) FROM search").fetchone()[0], 1)
        finally:
            connection.close()

    def test_missing_source_and_wrong_operation_fail_closed(self):
        connection = database()
        try:
            apply_job(connection, 1)
            self.assertTrue(apply_job(connection, 1, "delete").startswith("ignored"))
            with connection:
                connection.execute("DELETE FROM sources")
            self.assertEqual(retrieve(connection, "handoff")["context"], [])
            self.assertEqual(apply_job(connection, 1), "ignored: source missing")
        finally:
            connection.close()

    def test_literal_query_cannot_add_fts_operators(self):
        self.assertEqual(timeline("handoff OR nonexistent")["stages"][0]["matches"], [])
        self.assertEqual(timeline("")["stages"][0]["matches"], [])
        self.assertEqual(timeline('"; DROP TABLE sources; --')["stages"][0]["source"]["id"], "P-17")


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("command", choices=["run", "export", "test"])
    parser.add_argument("--stage", choices=[stage[0] for stage in STAGES], default="corrected")
    parser.add_argument("--query", default="handoff", help="One literal search phrase; default: handoff")
    args = parser.parse_args()
    if args.command == "test":
        suite = unittest.defaultTestLoader.loadTestsFromTestCase(ProofTests)
        raise SystemExit(not unittest.TextTestRunner(verbosity=2).run(suite).wasSuccessful())
    result = timeline(args.query)
    if args.command == "run":
        result = next(stage for stage in result["stages"] if stage["id"] == args.stage)
    print(json.dumps(result, indent=2))


if __name__ == "__main__":
    main()