"""Original LockedIn Labs recovery workshop. Python 3.11+ standard library.

Actually back up and restore a scratch application database while a separate
scratch destination retains its committed effects. No network, external service,
real business action, cloud restore, process crash or production RTO measurement.
"""
import argparse
from contextlib import closing
import json
from pathlib import Path
import sqlite3
import tempfile

CASES = ("unguarded", "retained-receipt", "new-operation-id", "pruned-receipt", "changed-payload")


def app_state(db):
    row = db.execute("SELECT job_id, operation_id, units, state, result_units FROM jobs").fetchone()
    return dict(zip(("job_id", "operation_id", "units", "state", "result_units"), row))


def destination_state(db):
    effects = [dict(zip(("sequence", "operation_id", "units"), row)) for row in
               db.execute("SELECT sequence, operation_id, units FROM effects ORDER BY sequence")]
    receipts = [dict(zip(("operation_id", "units", "result_units"), row)) for row in
                db.execute("SELECT operation_id, units, result_units FROM receipts ORDER BY operation_id")]
    return {"units": db.execute("SELECT coalesce(sum(units), 0) FROM effects").fetchone()[0],
            "effect_count": len(effects), "receipt_count": len(receipts),
            "effects": effects, "receipts": receipts}


def receive(db, operation_id, units, guarded):
    """One local destination transaction; it has no transaction with the app."""
    db.execute("BEGIN IMMEDIATE")
    try:
        prior = db.execute("SELECT units, result_units FROM receipts WHERE operation_id=?", (operation_id,)).fetchone() if guarded else None
        if prior is not None:
            if prior[0] != units:
                db.rollback()
                return {"state": "payload_conflict", "result_units": None}
            db.commit()
            return {"state": "replayed", "result_units": prior[1]}
        db.execute("INSERT INTO effects(operation_id, units) VALUES (?, ?)", (operation_id, units))
        result = db.execute("SELECT sum(units) FROM effects").fetchone()[0]
        if guarded:
            db.execute("INSERT INTO receipts VALUES (?, ?, ?)", (operation_id, units, result))
        db.commit()
        return {"state": "applied", "result_units": result}
    except BaseException:
        db.rollback()
        raise


def run_case(case, units=4):
    if case not in CASES:
        raise ValueError("unsupported case")
    if type(units) is not int or not 1 <= units <= 9:
        raise ValueError("units must be an integer from 1 through 9")
    guarded = case != "unguarded"
    with tempfile.TemporaryDirectory(prefix="fde-restore-replay-") as folder:
        base = Path(folder)
        with closing(sqlite3.connect(base / "app.db", isolation_level=None)) as app, closing(sqlite3.connect(base / "destination.db", isolation_level=None)) as dest:
            app.execute("CREATE TABLE jobs(job_id TEXT PRIMARY KEY, operation_id TEXT, units INTEGER, state TEXT, result_units INTEGER)")
            app.execute("INSERT INTO jobs VALUES ('JOB-1', 'OP-1', ?, 'pending', NULL)", (units,))
            dest.execute("CREATE TABLE effects(sequence INTEGER PRIMARY KEY, operation_id TEXT, units INTEGER)")
            dest.execute("CREATE TABLE receipts(operation_id TEXT PRIMARY KEY, units INTEGER, result_units INTEGER)")
            with closing(sqlite3.connect(base / "backup.db")) as backup:
                app.backup(backup)
                backup_job = app_state(backup)
            original_request = {"operation_id": "OP-1", "units": units}
            initial = receive(dest, **original_request, guarded=guarded)
            app.execute("UPDATE jobs SET state='completed', result_units=? WHERE job_id='JOB-1'", (initial["result_units"],))
            original_completed = app_state(app)
            before_restore = destination_state(dest)
            with closing(sqlite3.connect(base / "backup.db")) as backup, closing(sqlite3.connect(base / "restored.db", isolation_level=None)) as restored:
                backup.backup(restored)
                integrity = restored.execute("PRAGMA integrity_check").fetchone()[0]
                restored_pending = app_state(restored)  # Executes the app's actual read query.
                after_restore = destination_state(dest)
                pruned = 0
                if case == "pruned-receipt":
                    pruned = dest.execute("DELETE FROM receipts WHERE operation_id='OP-1'").rowcount
                before_retry = destination_state(dest)
                retry_request = {"operation_id": "OP-1-recovered" if case == "new-operation-id" else restored_pending["operation_id"],
                                 "units": units + 1 if case == "changed-payload" else restored_pending["units"]}
                retry = receive(dest, **retry_request, guarded=guarded)
                if retry["state"] in ("applied", "replayed"):
                    restored.execute("UPDATE jobs SET state='completed', result_units=? WHERE job_id='JOB-1'", (retry["result_units"],))
                return {"case": case, "original_request": original_request,
                        "backup_job": backup_job, "original_result": initial,
                        "original_application_after_completion": original_completed,
                        "destination_before_app_restore": before_restore,
                        "restored_application_before_retry": restored_pending,
                        "restored_integrity": integrity, "application_read_succeeded": True,
                        "destination_after_app_restore": after_restore,
                        "receipt_pruning": {"deleted_rows": pruned, "automatic_time_expiry_tested": False},
                        "destination_before_retry": before_retry, "retry_request": retry_request,
                        "retry_result": retry, "final_restored_application": app_state(restored),
                        "original_application_retained": app_state(app),
                        "final_destination": destination_state(dest),
                        "destination_integrity": dest.execute("PRAGMA integrity_check").fetchone()[0]}


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--case", choices=CASES, help="One comparison; default runs all five")
    parser.add_argument("--units", type=int, default=4, help="Original synthetic units, 1 through 9; default 4")
    args = parser.parse_args()
    if not 1 <= args.units <= 9:
        parser.error("--units must be from 1 through 9")
    report = {"schema": "lockedin.restore-replay-lab/v1", "sqlite_version": sqlite3.sqlite_version,
              "claim": "Actual local application backup and restore; separate destination retained. Sequential synthetic transactions, not a distributed recovery or production RTO proof.",
              "cases": [run_case(case, args.units) for case in ([args.case] if args.case else CASES)]}
    print(json.dumps(report, indent=2))


if __name__ == "__main__":
    main()
