#!/usr/bin/env python3
"""Original offline lab: queue completion, destination fencing and operation receipts.
Two owned processes, two scratch SQLite files, injected logical ticks. No network.
"""
import argparse
from contextlib import closing, contextmanager
import json
import multiprocessing as mp
import os
from pathlib import Path
import sqlite3
import sys
import tempfile

SCHEMA = "lockedin.worker-resource-fencing/v1"
CASES = {
    "finish-only": ("none", False, "none"),
    "fence-new-first": ("fence", False, "none"),
    "fence-old-first": ("fence", True, "none"),
    "same-fence-repeat": ("fence", False, "same"),
    "fence-and-receipt": ("receipt", True, "same"),
    "changed-operation": ("receipt", False, "changed"),
    "fence-rollback": ("receipt", False, "rollback"),
}


def connect(path):
    db = sqlite3.connect(path, isolation_level=None, timeout=3)
    db.row_factory = sqlite3.Row
    return db


@contextmanager
def transaction(db):
    db.execute("BEGIN IMMEDIATE")
    try:
        yield
    except BaseException:
        db.rollback()
        raise
    else:
        db.commit()


def initialize(queue, resource):
    with closing(connect(queue)) as db:
        db.executescript("CREATE TABLE job(id TEXT PRIMARY KEY, owner TEXT, fence INTEGER, lease_until INTEGER, attempts INTEGER, state TEXT); INSERT INTO job VALUES('JOB-1',NULL,0,0,0,'queued');")
    with closing(connect(resource)) as db:
        db.executescript("CREATE TABLE resource(id TEXT PRIMARY KEY, highest_fence INTEGER, units INTEGER); INSERT INTO resource VALUES('RESOURCE-1',0,0); CREATE TABLE receipt(operation TEXT PRIMARY KEY, units INTEGER NOT NULL);")


def job_state(path):
    with closing(connect(path)) as db:
        return dict(db.execute("SELECT * FROM job WHERE id='JOB-1'").fetchone())


def claim(path, owner, tick):
    with closing(connect(path)) as db, transaction(db):
        row = db.execute("SELECT * FROM job WHERE id='JOB-1'").fetchone()
        if row["state"] == "finished" or (row["state"] == "running" and tick < row["lease_until"]):
            raise ValueError("lease_busy_or_finished")
        fence = row["fence"] + 1
        db.execute("UPDATE job SET owner=?,fence=?,lease_until=?,attempts=attempts+1,state='running' WHERE id='JOB-1'", (owner, fence, tick + 5))
        return fence


def finish(path, owner, fence, tick):
    with closing(connect(path)) as db, transaction(db):
        updated = db.execute("UPDATE job SET state='finished' WHERE id='JOB-1' AND state='running' AND owner=? AND fence=? AND lease_until>?", (owner, fence, tick)).rowcount
        return "finished" if updated == 1 else "lease_lost"


def resource_state(path):
    with closing(connect(path)) as db:
        row = dict(db.execute("SELECT * FROM resource WHERE id='RESOURCE-1'").fetchone())
        row["receipt_count"] = db.execute("SELECT COUNT(*) FROM receipt").fetchone()[0]
        return row


class InjectedFailure(Exception):
    pass


def receive(path, mode, fence, units, operation="OP-1", fail_after_effect=False):
    if mode not in ("none", "fence", "receipt") or type(fence) is not int or fence < 1 or type(units) is not int or not 1 <= units <= 9 or operation != "OP-1":
        raise ValueError("invalid_fixture_input")
    before = resource_state(path)
    pending = None
    try:
        with closing(connect(path)) as db, transaction(db):
            row = db.execute("SELECT * FROM resource WHERE id='RESOURCE-1'").fetchone()
            if mode != "none" and fence < row["highest_fence"]:
                state = "fenced"
            else:
                if mode != "none":
                    db.execute("UPDATE resource SET highest_fence=MAX(highest_fence,?) WHERE id='RESOURCE-1'", (fence,))
                prior = db.execute("SELECT units FROM receipt WHERE operation=?", (operation,)).fetchone() if mode == "receipt" else None
                if prior and prior["units"] != units:
                    raise ValueError("operation_conflict")
                if prior:
                    state = "replayed"
                else:
                    db.execute("UPDATE resource SET units=units+? WHERE id='RESOURCE-1'", (units,))
                    if fail_after_effect:
                        # Only this failing connection sees these uncommitted changes.
                        pending = dict(db.execute("SELECT * FROM resource WHERE id='RESOURCE-1'").fetchone())
                        pending["receipt_count"] = db.execute("SELECT COUNT(*) FROM receipt").fetchone()[0]
                        raise InjectedFailure()
                    if mode == "receipt":
                        db.execute("INSERT INTO receipt VALUES(?,?)", (operation, units))
                    state = "applied"
    except InjectedFailure:
        state = "rolled_back"
    except ValueError as error:
        if str(error) != "operation_conflict":
            raise
        state = "operation_conflict"
    result = {"fence": fence, "operation": operation, "units": units, "state": state, "before": before, "after": resource_state(path)}
    if pending is not None:
        result["uncommitted_failure_view"] = pending
    return result


def old_worker(queue, resource, mode, units, pipe):
    try:
        fence = claim(queue, "worker-a", 0)
        # A real precheck before the pause. No transaction spans the two databases.
        current = job_state(queue)
        pipe.send({"event": "checked_before_pause", "pid": os.getpid(), "fence": fence, "job": current})
        if not pipe.poll(8) or pipe.recv() != "continue":
            raise RuntimeError("release_not_received")
        result = receive(resource, mode, fence, units)
        completion = finish(queue, "worker-a", fence, 6)
        pipe.send({"actor": "worker-a", "pid": os.getpid(), "destination": result, "queue_finish": completion})
    finally:
        pipe.close()


def message(pipe):
    if not pipe.poll(6):
        raise RuntimeError("owned_worker_message_timeout")
    return pipe.recv()


def run_case(name, units):
    mode, old_first, repeat = CASES[name]
    with tempfile.TemporaryDirectory(prefix="lockedin-fencing-") as folder:
        queue, resource = str(Path(folder) / "queue.sqlite"), str(Path(folder) / "resource.sqlite")
        initialize(queue, resource)
        ctx = mp.get_context("spawn")
        parent, child = ctx.Pipe()
        process = ctx.Process(target=old_worker, args=(queue, resource, mode, units, child))
        process.start()
        child.close()
        try:
            ready = message(parent)
            fence = claim(queue, "worker-b", 5)
            reclaimed = job_state(queue)
            attempts = []
            if old_first:
                parent.send("continue")
                attempts.append(message(parent))
            attempts.append({"actor": "worker-b", "pid": os.getpid(), "destination": receive(resource, mode, fence, units, fail_after_effect=repeat == "rollback")})
            if repeat == "rollback":
                attempts[-1]["queue_after_destination"] = job_state(queue)
                parent.send("continue")
                attempts.append(message(parent))
                attempts.append({"actor": "worker-b", "pid": os.getpid(), "destination": receive(resource, mode, fence, units)})
            elif repeat != "none":
                changed = (units - 1 if units == 9 else units + 1) if repeat == "changed" else units
                attempts.append({"actor": "worker-b", "pid": os.getpid(), "destination": receive(resource, mode, fence, changed)})
            completed = finish(queue, "worker-b", fence, 6)
            if not old_first and repeat != "rollback":
                parent.send("continue")
                attempts.append(message(parent))
            process.join(5)
            if process.is_alive() or process.exitcode != 0:
                raise RuntimeError("owned_worker_did_not_finish")
            return {"case": name, "mode": mode, "old_first_at_destination": old_first, "precheck_before_pause": ready,
                    "after_reclaim": reclaimed, "attempts": attempts, "current_queue_finish": completed,
                    "final_job": job_state(queue), "final_destination": resource_state(resource), "old_worker_exit": process.exitcode}
        finally:
            parent.close()
            if process.is_alive():
                process.terminate()
                process.join(2)
            if process.is_alive():
                process.kill()
                process.join(2)
            process.close()


def bounded_units(value):
    try:
        n = int(value)
    except ValueError as error:
        raise argparse.ArgumentTypeError("units must be an integer from 1 through 9") from error
    if str(n) != value or not 1 <= n <= 9:
        raise argparse.ArgumentTypeError("units must be an integer from 1 through 9")
    return n


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--case", choices=["all", *CASES], default="all")
    parser.add_argument("--units", type=bounded_units, default=4)
    args = parser.parse_args()
    selected = list(CASES) if args.case == "all" else [args.case]
    report = {"schema": SCHEMA, "fence_contract": "highest_fence is the highest token committed by an admitted, non-conflicting destination transaction; a rolled-back attempt does not advance it.", "runtime": {"python": sys.version.split()[0], "sqlite": sqlite3.sqlite_version},
              "logical_ticks": {"old_claim": 0, "new_claim": 5, "receiving_and_completion": 6},
              "scope": "Two owned local processes and two temporary databases; trusted synthetic tokens, no network or authentication.",
              "cases": [run_case(name, args.units) for name in selected]}
    print(json.dumps(report, indent=2))


if __name__ == "__main__":
    mp.freeze_support()
    main()