#!/usr/bin/env python3
"""Apply a reviewed access plan only to the local acceptance database."""
from __future__ import annotations

import argparse
import datetime
import json
from pathlib import Path

from apply_weknora_identities import sql_run
from weknora_access_etl import ACCESS_TABLES, LEDGER, SOURCE, canonical, migration_sql
from weknora_category_etl import save_json
from weknora_identity_etl import TARGET_DATABASE, TARGET_TENANT


def snapshot():
    fields = ["'database',current_database()",
              """'users',(SELECT coalesce(json_agg(json_build_object('id',id,'fingerprint',md5(to_jsonb(u)::text)) ORDER BY id),'[]') FROM users u)""",
              """'memberships',(SELECT coalesce(json_agg(json_build_object('id',id,'fingerprint',md5(to_jsonb(m)::text)) ORDER BY id),'[]') FROM tenant_members m)"""]
    fields += [f"'{table}',(SELECT coalesce(json_agg(to_jsonb(a) ORDER BY to_jsonb(a)::text),'[]') FROM {table} a)" for table in ACCESS_TABLES]
    result = json.loads(sql_run("BEGIN ISOLATION LEVEL REPEATABLE READ READ ONLY; SELECT json_build_object(" + ",".join(fields) + "); COMMIT;"))
    exists = sql_run(f"SELECT to_regclass('{LEDGER}') IS NOT NULL;") == "t"
    result[LEDGER] = json.loads(sql_run(f"SELECT coalesce(json_agg(to_jsonb(a) ORDER BY entity_key),'[]') FROM {LEDGER} a;")) if exists else []
    return result


def verification(plan, before, after):
    checks = {"acceptance_scope": before["database"] == after["database"] == TARGET_DATABASE,
              "all_users_unchanged": before["users"] == after["users"],
              "all_memberships_unchanged": before["memberships"] == after["memberships"],
              "all_explicit_grants_unchanged": before[ACCESS_TABLES[3]] == after[ACCESS_TABLES[3]]}
    expected = {
        ACCESS_TABLES[0]: {r["id"] for r in plan["roles"]},
        ACCESS_TABLES[1]: {(TARGET_TENANT, a["user_id"]) for a in plan["assignments"]},
        ACCESS_TABLES[2]: {(TARGET_TENANT, p["kind"], p["resource_id"]) for p in plan["publications"]},
    }
    def key(table, row):
        if table == ACCESS_TABLES[0]:
            return row["id"]
        return (row["tenant_id"], row["user_id"]) if table == ACCESS_TABLES[1] else (row["tenant_id"], row["kind"], row["resource_id"])
    for table in expected:
        old, new = ({key(table, r): r for r in state[table]} for state in (before, after))
        checks[table + "_existing_unchanged"] = all(new.get(k) == r for k, r in old.items())
        checks[table + "_planned_only"] = set(new) - set(old) == expected[table] - set(old)
    saved = {r["entity_key"]: r["source_spec"] for r in after[LEDGER]
             if r["source_system"] == SOURCE and r["target_tenant_id"] == TARGET_TENANT}
    planned = {"role/" + r["source_role"]: r for r in plan["roles"]}
    planned.update({"assignment/" + a["user_id"]: a for a in plan["assignments"]})
    planned.update({"publication/" + p["kind"] + "/" + p["resource_id"]: p for p in plan["publications"]})
    checks["durable_provenance"] = all(saved.get(k) == v for k, v in planned.items())
    return {"passed": all(checks.values()), "checks": checks}


def rollback_sql(plan, before, after, commit=False):
    migration_sql(plan)  # Reuse the scope and plan-integrity validation.
    inserted = {table: [r for r in after[table] if r not in before[table]] for table in (*ACCESS_TABLES, LEDGER)}
    payload = "'" + canonical(inserted).replace("'", "''") + "'::jsonb"
    parts = [f"""-- Removes only unchanged rows inserted by this reviewed application.
-- Default rehearsal. Emit COMMIT only after reviewing the current native edits.
BEGIN ISOLATION LEVEL SERIALIZABLE;
SET LOCAL statement_timeout='30s'; SET LOCAL lock_timeout='5s';
DO $scope$ BEGIN
 IF current_database()<>'{TARGET_DATABASE}' THEN RAISE EXCEPTION 'Wrong rollback database'; END IF;
END $scope$;
SELECT pg_advisory_xact_lock({TARGET_TENANT},11905);
LOCK TABLE users,tenant_members,workbench_access_roles,workbench_access_assignments,
 workbench_access_publications,workbench_access_grants,{LEDGER} IN SHARE ROW EXCLUSIVE MODE;
DO $suspended$ BEGIN
 IF EXISTS(SELECT 1 FROM users u JOIN yawan_legacy_identity_mappings l ON l.target_user_id=u.id
   JOIN tenant_members m ON m.user_id=u.id
   WHERE l.target_tenant_id={TARGET_TENANT} AND l.source_user_id<>1
     AND (u.is_active OR m.status<>'suspended' OR m.role<>'contributor' OR m.tenant_id<>{TARGET_TENANT}))
 THEN RAISE EXCEPTION 'Account activation changed; review rollback policy first'; END IF;
END $suspended$;
CREATE TEMP TABLE _access_rollback_payload ON COMMIT DROP AS SELECT {payload} AS data;"""]
    for table in (*ACCESS_TABLES, LEDGER):
        parts.append(f"""DO $unchanged$ BEGIN
 IF EXISTS(SELECT 1 FROM _access_rollback_payload p, jsonb_array_elements(p.data->'{table}') expected
   WHERE NOT EXISTS(SELECT 1 FROM {table} actual WHERE to_jsonb(actual)=expected))
 THEN RAISE EXCEPTION 'Native edits or missing rows in {table}; rollback aborted'; END IF;
END $unchanged$;""")
    parts.append("""DO $references$ BEGIN
 IF EXISTS(SELECT 1 FROM workbench_access_grants g JOIN workbench_access_publications p USING(tenant_id,kind,resource_id)
   WHERE to_jsonb(p) IN (SELECT value FROM _access_rollback_payload, jsonb_array_elements(data->'workbench_access_publications')))
 OR EXISTS(SELECT 1 FROM workbench_access_assignments a JOIN workbench_access_roles r ON r.id=a.role_id
   WHERE to_jsonb(r) IN (SELECT value FROM _access_rollback_payload, jsonb_array_elements(data->'workbench_access_roles'))
     AND to_jsonb(a) NOT IN (SELECT value FROM _access_rollback_payload, jsonb_array_elements(data->'workbench_access_assignments')))
 THEN RAISE EXCEPTION 'New references to migrated access rows; rollback aborted'; END IF;
END $references$;""")
    for table in (ACCESS_TABLES[1], ACCESS_TABLES[2], ACCESS_TABLES[0], LEDGER):
        parts.append(f"DELETE FROM {table} actual WHERE to_jsonb(actual) IN (SELECT value FROM _access_rollback_payload,jsonb_array_elements(data->'{table}'));")
    parts.append("COMMIT;" if commit else "ROLLBACK;")
    return "\n".join(parts) + "\n"


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("command", choices=("plan", "apply"))
    parser.add_argument("--manifest", type=Path, required=True)
    parser.add_argument("--out", type=Path, required=True)
    args = parser.parse_args()
    suffixes = (".before.json", ".after.json", ".attempt.json", ".committed.json", ".rollback.sql")
    if args.out.exists() or any(args.out.with_suffix(s).exists() for s in suffixes):
        parser.error("Checkpoint exists; choose a new output name")
    plan = json.loads(args.manifest.read_text())
    sql = migration_sql(plan)
    before = snapshot()
    save_json(args.out.with_suffix(".before.json"), before)
    rehearsal = [json.loads(line) for line in sql_run(sql).splitlines() if line.startswith("{")]
    if snapshot() != before:
        raise RuntimeError("Rehearsal changed state or concurrent changes detected")
    result = {"checked_at": datetime.datetime.now(datetime.timezone.utc).isoformat(), "database": TARGET_DATABASE,
              "tenant_id": TARGET_TENANT, "manifest": str(args.manifest), "plan_sha256": plan["plan_sha256"],
              "rehearsal": rehearsal, "committed": False}
    if args.command == "apply":
        save_json(args.out.with_suffix(".attempt.json"), {**result, "commit_requested": True,
            "recovery": "Reconcile durable access ledger and checkpoints; never recreate records blindly"})
        changes = [json.loads(line) for line in sql_run(migration_sql(plan, True)).splitlines() if line.startswith("{")]
        result.update(committed=True, changes=changes)
        save_json(args.out.with_suffix(".committed.json"), result)
        after = snapshot()
        save_json(args.out.with_suffix(".after.json"), after)
        result["verification"] = verification(plan, before, after)
        args.out.with_suffix(".rollback.sql").write_text(rollback_sql(plan, before, after), encoding="utf-8")
    save_json(args.out, result)
    print(json.dumps(result))
    return 0 if result.get("verification", {}).get("passed", True) else 1


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