#!/usr/bin/env python3
from __future__ import annotations

import argparse
import os
import subprocess
import sys

import psycopg


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Run a command while holding a Postgres advisory lock.",
    )
    parser.add_argument("--dsn", required=True, help="Postgres DSN")
    parser.add_argument("--key", required=True, type=int, help="Advisory lock key")
    parser.add_argument("--label", required=True, help="Human-readable lock label")
    parser.add_argument("command", nargs=argparse.REMAINDER, help="Command to run after --")
    return parser.parse_args()


def main() -> int:
    args = parse_args()
    command = list(args.command)
    if command and command[0] == "--":
        command = command[1:]
    if not command:
        print("No command provided for advisory lock wrapper.", file=sys.stderr)
        return 2

    with psycopg.connect(args.dsn, autocommit=True) as conn:
        with conn.cursor() as cur:
            cur.execute("select pg_try_advisory_lock(%s)", (args.key,))
            acquired = cur.fetchone()
            if not acquired or not acquired[0]:
                print(f"{args.label} already running; skipping.", file=sys.stderr)
                return 0

        env = os.environ.copy()
        env["UFORIA_REFRESH_LOCK_HELD"] = "1"
        try:
            result = subprocess.run(command, env=env, check=False)
            return result.returncode
        finally:
            with conn.cursor() as cur:
                cur.execute("select pg_advisory_unlock(%s)", (args.key,))


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