#!/usr/bin/env python3
"""Original PostgreSQL 17 lock-queue lab. Python 3.11+, Docker, no pip packages.

Keep fde-postgres-row-security.py beside this file: its Lab owns the disposable
container lifecycle. Only its lifecycle and admin helpers are reused. No image
pull, host ports, network, existing database or application credentials. Local
trust supplies fictional roles. Three real psql sessions and one observer run
ordered interleavings; this is not a workload or migration-duration benchmark.
"""
import argparse
import importlib.util
import json
from pathlib import Path
import queue
import subprocess
import sys
import threading
import time
import uuid


def runtime():
    companion = Path(__file__).resolve().with_name("fde-postgres-row-security.py")
    if not companion.is_file():
        raise RuntimeError("Keep the downloaded fde-postgres-row-security.py companion beside this file")
    spec = importlib.util.spec_from_file_location("fde_owned_postgres", companion)
    module = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(module)
    return module


class Session:
    """One persistent psql connection; output markers delimit completed batches."""
    def __init__(self, container, role):
        self.role = role
        self.commands = []
        self.output = queue.Queue()
        self.diagnostics = []
        self.process = subprocess.Popen([
            "docker", "exec", "-i", container, "psql", "-X", "-w", "-qAt",
            "--username", role, "--dbname", "postgres", "--set", "ON_ERROR_STOP=0",
            "--set", "VERBOSITY=sqlstate"], stdin=subprocess.PIPE,
            stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True, bufsize=1)
        def stdout():
            for line in self.process.stdout:
                self.output.put(line.rstrip("\n"))
            self.output.put(None)
        def stderr():
            self.diagnostics.extend(line.rstrip("\n") for line in self.process.stderr)
        self.threads = [threading.Thread(target=stdout, daemon=True), threading.Thread(target=stderr, daemon=True)]
        for thread in self.threads:
            thread.start()
        self.pid = int(self.finish(self.send("SELECT pg_backend_pid();"))[0])

    def send(self, sql):
        marker = "FDE_END_" + uuid.uuid4().hex
        self.commands.append(sql)
        self.process.stdin.write(sql + "\n\\echo " + marker + "\n")
        self.process.stdin.flush()
        return marker

    def finish(self, marker):
        lines = []
        deadline = time.monotonic() + 22
        while True:
            line = self.output.get(timeout=max(0.01, deadline - time.monotonic()))
            if line == marker:
                return lines
            if line is None:
                raise RuntimeError("Owned psql session ended before its batch marker")
            lines.append(line)

    def close(self):
        if self.process.poll() is None:
            try:
                self.process.stdin.write("\\q\n")
                self.process.stdin.flush()
                self.process.wait(timeout=3)
            except (BrokenPipeError, subprocess.TimeoutExpired):
                self.process.kill()
                self.process.wait(timeout=3)
        for thread in self.threads:
            thread.join(timeout=2)
        for stream in (self.process.stdin, self.process.stdout, self.process.stderr):
            stream.close()
        return {"pid": self.pid, "role": self.role, "commands": self.commands,
                "diagnostic": "\n".join(self.diagnostics), "psql_exit_code": self.process.returncode}


SETUP = """
CREATE ROLE lab_reader LOGIN NOSUPERUSER NOCREATEDB NOCREATEROLE NOBYPASSRLS;
CREATE ROLE lab_migrator LOGIN NOSUPERUSER NOCREATEDB NOCREATEROLE NOBYPASSRLS;
CREATE SCHEMA lab AUTHORIZATION lab_migrator;
SET ROLE lab_migrator;
CREATE TABLE lab.jobs (id integer PRIMARY KEY, owner_name text NOT NULL);
INSERT INTO lab.jobs VALUES (42, 'Queue A');
GRANT USAGE ON SCHEMA lab TO lab_reader;
GRANT SELECT ON lab.jobs TO lab_reader;
RESET ROLE;
"""
READ = "SELECT json_build_object('pid', pg_backend_pid(), 'role', current_user, 'rows', (SELECT json_agg(j ORDER BY id) FROM (SELECT id, owner_name FROM lab.jobs) j));"
ALTER = "ALTER TABLE lab.jobs ADD COLUMN owner_display_name text;"


def snapshot(lab, sessions):
    pids = ",".join(str(session.pid) for session in sessions)
    sql = """SELECT json_build_object(
      'sessions', (SELECT json_agg(a ORDER BY a.pid) FROM
        (SELECT pid, usename, state, wait_event_type, wait_event,
          pg_blocking_pids(pid) AS blockers FROM pg_stat_activity WHERE pid IN (%s)) a),
      'table_locks', coalesce((SELECT json_agg(l ORDER BY l.pid, l.mode) FROM
        (SELECT pid, mode, granted FROM pg_locks WHERE relation = 'lab.jobs'::regclass
          AND pid IN (%s)) l), '[]'::json));""" % (pids, pids)
    return {"sql": sql, **json.loads(lab.admin(sql)["output"])}


def wait_queued(lab, sessions, target, mode):
    # Observe a real wait before releasing another session; elapsed time is not
    # the coordination signal. A missing wait fails this run rather than passes.
    deadline = time.monotonic() + 4
    while time.monotonic() < deadline:
        observed = snapshot(lab, sessions)
        if any(row["pid"] == target.pid and row["mode"] == mode and not row["granted"]
               for row in observed["table_locks"]):
            return observed
        time.sleep(0.04)
    raise RuntimeError("Expected owned lock wait was not observed; no completed-case claim")


def run_case(lab, case, timeout_ms, reader_end):
    sessions = []
    transcripts = []
    try:
        blocker = Session(lab.container_id, "lab_reader"); sessions.append(blocker)
        ddl = Session(lab.container_id, "lab_migrator"); sessions.append(ddl)
        follower = Session(lab.container_id, "lab_reader"); sessions.append(follower)
        first = blocker.finish(blocker.send("BEGIN; SET LOCAL idle_in_transaction_session_timeout = '25s';\n" + READ))
        before = snapshot(lab, sessions)
        # All DDL attempts are bounded. Cases controlled by cancel/commit use a
        # larger guard so the authored action, rather than timeout, resolves them.
        budget = timeout_ms if case == "lock-timeout" else 15000
        ddl_sql = "BEGIN; SET LOCAL lock_timeout = '%dms'; SET LOCAL statement_timeout = '20s';\n%s\n\\echo FDE_DDL_SQLSTATE :SQLSTATE\nCOMMIT;" % (budget, ALTER)
        ddl_marker = ddl.send(ddl_sql)
        queued_ddl = wait_queued(lab, sessions, ddl, "AccessExclusiveLock")
        follower_marker = follower.send("SET statement_timeout = '20s';\n" + READ)
        queued_read = wait_queued(lab, sessions, follower, "AccessShareLock")
        action = None
        if case == "cancel-queued-ddl":
            action = lab.admin("SELECT pg_cancel_backend(%d);" % ddl.pid)
        elif case == "release-reader":
            end_sql = reader_end.upper() + ";"
            action = {"sql": end_sql, "output": blocker.finish(blocker.send(end_sql))}
        ddl_output = ddl.finish(ddl_marker)
        follower_output = follower.finish(follower_marker)
        state_line = [line for line in ddl_output if line.startswith("FDE_DDL_SQLSTATE ")]
        if len(state_line) != 1:
            raise RuntimeError("DDL completion SQLSTATE was not retained")
        ddl_state = state_line[0].split()[-1]
        reads = [json.loads(line) for line in follower_output if line.startswith("{")]
        if len(reads) != 1:
            raise RuntimeError("Later read did not produce exactly one observed result")
        after = snapshot(lab, sessions)
        columns_sql = "SELECT json_agg(column_name ORDER BY ordinal_position) FROM information_schema.columns WHERE table_schema = 'lab' AND table_name = 'jobs';"
        columns = json.loads(lab.admin(columns_sql)["output"])
        result = {"case": case, "pids": {"blocker": blocker.pid, "ddl": ddl.pid, "follower": follower.pid},
                  "first_read": json.loads(first[0]), "before_ddl": before,
                  "ddl_queued": queued_ddl, "later_read_queued": queued_read,
                  "resolution_action": action, "ddl_sqlstate": ddl_state,
                  "ddl_output": ddl_output, "later_read": reads[0], "after_resolution": after,
                  "columns_sql": columns_sql, "columns": columns, "lock_timeout_ms": budget}
        blocker.finish(blocker.send("ROLLBACK;"))
        return result
    finally:
        # End only these owned synthetic backends before closing local clients.
        try:
            if sessions:
                lab.admin("SELECT pg_cancel_backend(pid) FROM pg_stat_activity WHERE pid IN (%s) AND state = 'active';" % ",".join(str(s.pid) for s in sessions))
        finally:
            for session in sessions:
                transcripts.append(session.close())
        if 'result' in locals():
            result["session_transcripts"] = transcripts


def run(timeout_ms, reader_end="commit"):
    support = runtime()
    lab = support.Lab()
    report = None
    try:
        lab.start()
        lab.admin(SETUP)
        container = json.loads(support.command(["docker", "inspect", lab.container_id]).stdout)[0]
        cases = []
        for case in ("cancel-queued-ddl", "lock-timeout", "release-reader"):
            cases.append(run_case(lab, case, timeout_ms, reader_end))
        report = {"scope": "Three ordered, actual PostgreSQL lock interleavings; no production load, authentication, pool or zero-downtime claim.",
                  "postgres_version": lab.admin("SELECT version();")["output"],
                  "container_id": lab.container_id, "image_id": lab.image["Id"],
                  "image_repo_digests": lab.image.get("RepoDigests", []),
                  "network_mode": container["HostConfig"]["NetworkMode"],
                  "host_ports_published": bool(container["HostConfig"].get("PortBindings")),
                  "identity_fixture": "Local trust; ordinary reader and migration owner; postgres observer can inspect/cancel only owned synthetic sessions.",
                  "setup_sql": SETUP, "reader_end": reader_end,
                  "role_attributes": json.loads(lab.admin("SELECT json_agg(r) FROM (SELECT rolname, rolsuper, rolbypassrls FROM pg_roles WHERE rolname IN ('lab_reader', 'lab_migrator') ORDER BY rolname) r;")["output"]),
                  "cases": cases}
    finally:
        removed = lab.remove()
        if report is not None:
            report["owned_container_removed"] = removed
        if not removed:
            raise RuntimeError("Owned scratch container cleanup was not confirmed: " + lab.name)
    return report


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--lock-timeout-ms", type=int, default=6000,
                        help="Timeout-case budget only: 5000 through 12000ms (default 6000)")
    parser.add_argument("--reader-end", choices=("commit", "rollback"), default="commit",
                        help="End the read-only blocker in the release-reader case (default commit)")
    args = parser.parse_args()
    if not 5000 <= args.lock_timeout_ms <= 12000:
        parser.error("--lock-timeout-ms must be from 5000 through 12000")
    try:
        print(json.dumps(run(args.lock_timeout_ms, args.reader_end), indent=2))
    except (RuntimeError, subprocess.SubprocessError, queue.Empty, OSError, ValueError) as error:
        print("Lab incomplete: " + str(error), file=sys.stderr)
        return 1
    return 0


if __name__ == "__main__":
    raise SystemExit(main())
