#!/usr/bin/env python3
"""Apply checked migrations atomically from the trusted deployment runner."""
import argparse
import hashlib
from pathlib import Path
import re
import subprocess

parser = argparse.ArgumentParser()
parser.add_argument("--container", default="endoleap-dev-postgres-1")
parser.add_argument("--database", default="foundry_auth")
parser.add_argument("--role", default="foundry")
args = parser.parse_args()
for value in (args.database, args.role):
    if not re.fullmatch(r"[a-z_][a-z0-9_]{0,62}", value):
        raise SystemExit("Invalid deployment database or role")
migrations = sorted((Path(__file__).resolve().parent.parent / "db").glob("[0-9]*.sql"))
if not migrations:
    raise SystemExit("No migrations found")

script = [
    "\\set ON_ERROR_STOP on",
    "BEGIN;",
    "SET LOCAL lock_timeout='30s';",
    "SELECT pg_advisory_xact_lock(731920164);",
    "CREATE TABLE IF NOT EXISTS foundry_schema_migrations "
    "(name text PRIMARY KEY, checksum text NOT NULL, applied_at timestamptz NOT NULL DEFAULT now());",
]
tables = set()
for index, path in enumerate(migrations):
    source = path.read_text()
    name = path.name
    if not re.fullmatch(r"[0-9]+-[a-z0-9-]+\.sql", name):
        raise SystemExit("Invalid migration name")
    checksum = hashlib.sha256(source.encode()).hexdigest()
    tables.update(re.findall(r"CREATE TABLE ([a-z_]+)", source))
    # The first deployed version predates migration tracking.
    if index == 0:
        script.extend([
            "SELECT bool_and(to_regclass('public.' || name) IS NOT NULL) AS baseline FROM "
            "unnest(ARRAY['project_accounts','oauth_attempts','login_grants','project_sessions']) name \\gset",
            "\\if :baseline",
            f"INSERT INTO foundry_schema_migrations(name,checksum) VALUES('{name}','{checksum}') ON CONFLICT DO NOTHING;",
            "\\endif",
        ])
    script.extend([
        f"SELECT EXISTS(SELECT 1 FROM foundry_schema_migrations WHERE name='{name}') AS applied, "
        f"EXISTS(SELECT 1 FROM foundry_schema_migrations WHERE name='{name}' AND checksum='{checksum}') AS valid \\gset",
        "\\if :applied",
        "\\if :valid",
        "\\else",
        "\\echo Migration checksum mismatch",
        "\\quit 1",
        "\\endif",
        "\\else",
    ])
    source = source.strip()
    if not source.startswith("BEGIN;") or not source.endswith("COMMIT;"):
        raise SystemExit("Migration must have explicit transaction boundaries")
    script.append(source[len("BEGIN;"):-len("COMMIT;")])
    script.extend([
        f"INSERT INTO foundry_schema_migrations(name,checksum) VALUES('{name}','{checksum}');",
        "\\endif",
    ])
for table in sorted(tables):
    script.append(f'GRANT SELECT,INSERT,UPDATE,DELETE ON TABLE "{table}" TO "{args.role}";')
script.extend([
    f'GRANT USAGE,SELECT ON SEQUENCE link_clicks_id_seq TO "{args.role}";',
    "COMMIT;",
])
subprocess.run(
    ["docker", "exec", "-i", args.container, "sh", "-c",
     'exec psql -X -U "${POSTGRES_USER:-postgres}" -d "$1" -v ON_ERROR_STOP=1', "sh", args.database],
    input=("\n".join(script) + "\n").encode(), check=True,
)
print("Schema migrations and application grants verified.")