#!/usr/bin/env python3
"""Writes results/exp3_catalogue_counts.json and results/exp3_scope_history.json
from rebuilt schemas wc_at_016, wc_at_022 and wc_head_clean (see README)."""
import json, os, subprocess
HERE = os.path.dirname(os.path.abspath(__file__)); RES = os.path.join(HERE, "..", "results")
PSQL = os.environ.get("PSQL", "/opt/homebrew/opt/postgresql@17/bin/psql")
PORT = os.environ.get("PGPORT_EXP", "55499")

def q(db, sql):
    out = subprocess.run([PSQL, "-h", "localhost", "-p", PORT, "-U", "postgres", "-d", db, "-At", "-F", "\t", "-c", sql],
                         check=True, capture_output=True, text=True).stdout.strip()
    return [l.split("\t") for l in out.splitlines()]

COUNTS = """
with t as (
  select c.oid, c.relname, c.relrowsecurity rls, c.relforcerowsecurity forced,
         (select count(*) from pg_policy p where p.polrelid=c.oid) pol,
         (has_table_privilege('wayclub_app', c.oid, 'SELECT') or has_table_privilege('wayclub_app', c.oid, 'INSERT')
          or has_table_privilege('wayclub_app', c.oid, 'UPDATE') or has_table_privilege('wayclub_app', c.oid, 'DELETE')) app_grant
  from pg_class c join pg_namespace n on n.oid=c.relnamespace
  where n.nspname='public' and c.relkind='r' and c.relname <> 'schema_migrations')
select count(*), count(*) filter (where rls), count(*) filter (where not rls),
       coalesce(string_agg(relname, ', ') filter (where not rls), ''), count(*) filter (where forced),
       sum(pol), count(*) filter (where pol>0), count(*) filter (where app_grant),
       count(*) filter (where rls and pol=0 and not app_grant)
from t"""
ROLES = """select count(*) filter (where polroles && array[0::oid,'wayclub_app'::regrole::oid]),
       count(*) filter (where polroles && array['wayclub_service'::regrole::oid]),
       count(*) filter (where coalesce(pg_get_expr(polqual,polrelid),'')||coalesce(pg_get_expr(polwithcheck,polrelid),'') ~ 'app\\.'),
       count(*)
from pg_policy p join pg_class c on c.oid=p.polrelid join pg_namespace n on n.oid=c.relnamespace where n.nspname='public'"""
DEFINERS = """select count(*), count(*) filter (where proconfig is not null) from pg_proc p join pg_namespace n on n.oid=p.pronamespace
where p.prosecdef and n.nspname in ('app','wayclub_auth','public','abuse','platform')"""
SCOPE = open(os.path.join(HERE, "..", "sql", "old_vs_new_predicate.sql")).read()

counts = {}
for key, db in (("m016", "wc_at_016"), ("m022", "wc_at_022"), ("head", "wc_head_clean")):
    r = q(db, COUNTS)[0]
    d = dict(zip(["tables", "rls_enabled", "rls_disabled", "rls_disabled_names", "rls_forced", "policies",
                  "tables_with_policies", "app_reachable", "deny_all"], r))
    for k in d:
        if k != "rls_disabled_names":
            d[k] = int(d[k])
    ro = q(db, ROLES)[0]
    d.update(request_policies=int(ro[0]), service_policies=int(ro[1]), policies_calling_app_helpers=int(ro[2]), public_policies=int(ro[3]))
    de = q(db, DEFINERS)[0]
    d.update(definer_functions=int(de[0]), definers_with_set=int(de[1]))
    # build_db_at.sh applies files directly, so count the migration files at the commit.
    commit = {"m016": "85475d7", "m022": "940b2ae", "head": "65a18ff"}[key]
    repo = os.environ.get("WAYCLUB_REPO")
    d["migration_files"] = (sum(1 for f in subprocess.run(["git", "-C", repo, "ls-tree", "--name-only", commit, "infra/postgres/migrations/"],
                            check=True, capture_output=True, text=True).stdout.split() if f.endswith(".sql")) if repo else None)
    d["commit"] = commit
    d["database"] = db
    counts[key] = d
json.dump(counts, open(os.path.join(RES, "exp3_catalogue_counts.json"), "w"), indent=2)

hist = []
for label, db in (("016 (85475d7)", "wc_at_016"), ("022 (940b2ae)", "wc_at_022"), ("HEAD (65a18ff)", "wc_head_clean")):
    r = q(db, SCOPE.split("\n", 1)[1])[0]
    hist.append({"label": label, "database": db, "old_scope": int(r[0]), "new_scope": int(r[1]), "blind_spot": int(r[2]),
                 "blind_spot_tables": r[3], "blind_spot_without_rls": int(r[4])})
json.dump(hist, open(os.path.join(RES, "exp3_scope_history.json"), "w"), indent=2)
print(json.dumps(counts, indent=1)); print(json.dumps(hist, indent=1))