#!/usr/bin/env python3
"""Refresh staging from a production dump without losing a row a person wrote on staging.

Sources of truth: this script, `main/app/cli_refresh.py` (the export and apply halves, and the
policy for a row that conflicts), `main/app/services/discussion.py` (`CLIENT_WRITTEN_TABLES`),
`main/docs/RUNBOOK.md` (Staging refresh). Run on the host, from the workspace root, with the
system python (it drives docker; nothing here imports the app):

    python3 main/scripts/staging-refresh.py --from backups/production/<file>.dump --check
    python3 main/scripts/staging-refresh.py --from backups/production/<file>.dump

Adam comments on staging. A plain restore replaces `discussion_comments` and every other
table people write, and a lost comment is not recoverable, so this runs in a fixed order:

1. EXPORT every client-written row through `app.cli staging-refresh export` into
   `backups/refresh/<stamp>-client-rows.json`, and re-read the file before going on. This is
   the recovery artefact: whatever fails later, the rows are on disk. `--check` stops here,
   having restored nothing, and prints the counts and the commands it would run.
2. Drop and recreate the `public` schema in the database container, then `pg_restore` the
   dump into it. The whole schema goes, not `--clean` table by table: a staging-only table
   (one a migration deployed here first) would otherwise block the drop of a table it points at.
3. `alembic upgrade head` in the app container, so a dump from a production that is behind
   staging comes up to staging's schema, then the always-safe backfills.
4. APPLY the exported rows (`app.cli staging-refresh apply`): by natural key, never losing one;
   the report says what was restored, what was already there, what was renumbered.
5. Revoke every session (production's session rows are not staging's; sign in again).

The app container is not stopped; a request during the restore fails once and recovers. The
export file is kept. Nothing here touches production: the dump is read, never written.
"""

from __future__ import annotations

import argparse
import datetime as dt
import json
import pathlib
import shlex
import subprocess
import sys

ROOT = pathlib.Path(__file__).resolve().parents[2]
REFRESH_DIR = ROOT / "backups" / "refresh"
#: Backfills that are safe on any restored dump (each is idempotent and reports its work).
POST_RESTORE_BACKFILLS = ("accounts", "levels")


def sh(cmd: str, *, stdin=None, capture: bool = False, cwd: pathlib.Path | None = None, check: bool = True):
    print(f"$ {cmd}" + (" < …" if stdin is not None else ""), flush=True)
    return subprocess.run(shlex.split(cmd), stdin=stdin, cwd=cwd, check=check, text=True,
                          capture_output=capture)


def fail(msg: str, code: int = 1) -> None:
    print(f"staging-refresh: {msg}", file=sys.stderr)
    sys.exit(code)


def main() -> int:
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--from", dest="dump", required=True, help="a custom-format pg_dump of production")
    ap.add_argument("--check", action="store_true", help="export to the file and print the plan; restore nothing")
    ap.add_argument("--db", default="dfp", help="the database to refresh (default dfp)")
    ap.add_argument("--db-container", default="dutyfreeprofessor-db",
                    help="where psql and pg_restore run (default dutyfreeprofessor-db; the rehearsal uses dfp-devdb)")
    ap.add_argument("--app-exec", default="docker exec -i dutyfreeprofessor-app python -m app.cli",
                    help="how to run app.cli against the database being refreshed")
    ap.add_argument("--alembic-exec", default="docker exec -i dutyfreeprofessor-app alembic",
                    help="how to run alembic against the database being refreshed")
    ap.add_argument("--keep-sessions", action="store_true", help="do not revoke the restored sessions")
    ap.add_argument("--discard-decisions", action="store_true",
                    help="refresh even though the database holds catalogue decisions (aliases, merges, the ledger, pins), "
                         "which a refresh does not preserve")
    a = ap.parse_args()

    dump = pathlib.Path(a.dump)
    if not dump.is_file():
        fail(f"no such dump: {dump}")
    with dump.open("rb") as fh:
        if fh.read(5) != b"PGDMP":
            fail(f"{dump} is not a custom-format pg_dump (expected the PGDMP header)")
    app_cwd = ROOT / "main"  # a local app-exec runs from the instance; docker ignores cwd

    # 1. export, to disk, re-read
    REFRESH_DIR.mkdir(parents=True, exist_ok=True)
    stamp = dt.datetime.now(dt.UTC).strftime("%Y-%m-%d-%H%M")
    out = REFRESH_DIR / f"{stamp}-client-rows.json"
    exported = sh(f"{a.app_exec} staging-refresh export", capture=True, cwd=app_cwd)
    payload = json.loads(exported.stdout)
    tmp = out.with_suffix(".json.tmp")
    tmp.write_text(json.dumps(payload, indent=1), encoding="utf-8")
    tmp.replace(out)
    payload = json.loads(out.read_text(encoding="utf-8"))
    counts = payload.get("counts") or {}
    print(f"exported to {out.relative_to(ROOT)} (schema {payload.get('schema_version')}):")
    for name, n in counts.items():
        print(f"  {name}: {n}")
    if not counts:
        fail("the export names no tables; nothing was restored")

    dbx = f"docker exec -i {a.db_container}"
    # The catalogue's decisions are NOT in the preserved set (the catalogue decisions of 15 Sep,
    # §6): a refresh discards every alias, merge, ledger row, pin and ignore taken here. Count
    # them first and refuse unless told to discard them, so a desk session is never lost by
    # accident; the plan is to preserve them under natural keys.
    decided = sh(
        f"{dbx} psql -U dfp -d {a.db} -tA -v ON_ERROR_STOP=1 -c \"select 'decisions', count(*) from decisions "
        "union all select 'proposals', count(*) from proposals "
        "union all select 'aliases', count(*) from brands where alias_of_id is not null "
        "union all select 'line aliases', count(*) from product_lines where alias_of_id is not null "
        "union all select 'wordings decided', count(*) from attribute_aliases where decided_by is not null "
        "union all select 'pairs decided by a person', count(*) from suggestions where decided_by is not null "
        "union all select 'merges by a person', count(*) from merges where merged_by is not null "
        "union all select 'pins and ignores', count(*) from listings where pinned_variant_id is not null or ignored_at is not null\"",
        capture=True, check=False)
    if decided.returncode != 0:
        # A guard that cannot count must never let the refresh through: a dropped table once
        # would have made the refusal silently stop firing.
        fail(f"the decisions guard could not count (psql exit {decided.returncode}): {(decided.stderr or '').strip()[-500:]}", 4)
    held = {}
    for line in (decided.stdout or "").splitlines():
        if "|" in line:
            name, n = line.rsplit("|", 1)
            held[name.strip()] = int(n.strip() or 0)
    ledger_file = REFRESH_DIR / f"{stamp}-decisions.jsonl"
    if any(held.values()):
        print("catalogue decisions in the database being refreshed:")
        for name, n in held.items():
            if n:
                print(f"  {name}: {n}")
        if a.discard_decisions:
            print("  --discard-decisions: the ledger is NOT exported and will not come back")
        elif not a.check:
            # The ledger rides in the preserved set (K2, spec §8): exported by natural key before the
            # drop, replayed after the schema is up. Aliases, merges, pins and pair rulings are all
            # decision rows now, so this is the whole of it.
            exported_ledger = sh(f"{a.app_exec} decisions export --all --origin-host staging", capture=True, cwd=app_cwd)
            ledger_file.write_text(exported_ledger.stdout, encoding="utf-8")
            print(f"  exported the ledger to {ledger_file.relative_to(ROOT)}")
        else:
            print("  [check] the ledger would be exported before the drop and replayed after alembic upgrade head")
    if a.check:
        print("\n[check] nothing restored. The refresh would run, in order:")
        print(f"  {dbx} psql -U dfp -d {a.db} -v ON_ERROR_STOP=1 -c 'drop schema public cascade; create schema public;'")
        print(f"  {dbx} pg_restore -U dfp -d {a.db} --no-owner --no-privileges < {dump}")
        print(f"  {a.alembic_exec} upgrade head")
        for name in POST_RESTORE_BACKFILLS:
            print(f"  {a.app_exec} backfill {name}")
        print(f"  {a.app_exec} staging-refresh apply --from - < {out.relative_to(ROOT)}")
        if not a.keep_sessions:
            print(f"  {a.app_exec} sessions revoke --all")
        return 0

    # 2. the schema goes, the dump comes in
    sh(f"{dbx} psql -U dfp -d {a.db} -v ON_ERROR_STOP=1 -c 'drop schema public cascade; create schema public;'")
    with dump.open("rb") as fh:
        restored = sh(f"{dbx} pg_restore -U dfp -d {a.db} --no-owner --no-privileges", stdin=fh, capture=True, check=False)
    if restored.stderr.strip():
        print(restored.stderr.rstrip()[-2000:])
    version = sh(f"{dbx} psql -U dfp -d {a.db} -tA -c 'select version_num from alembic_version'",
                 capture=True, check=False)
    if version.returncode != 0 or not version.stdout.strip():
        fail(f"the restore left no alembic_version table; the rows are safe in {out.relative_to(ROOT)}")
    print(f"restored {dump.name}: schema {version.stdout.strip()}")

    # 3. up to staging's schema, then the always-safe backfills
    sh(f"{a.alembic_exec} upgrade head", cwd=app_cwd)
    for name in POST_RESTORE_BACKFILLS:
        sh(f"{a.app_exec} backfill {name}", cwd=app_cwd)

    # 4. the ledger comes back by natural key, then the client rows
    if ledger_file.exists() and not a.discard_decisions:
        replayed = sh(f"{a.app_exec} decisions replay --file {ledger_file} --label production-refresh --force", cwd=app_cwd, check=False)
        if replayed.returncode == 2:
            fail(f"the ledger replay refused; the export is safe in {ledger_file.relative_to(ROOT)}", 2)
        if replayed.returncode == 1:
            print(f"some decisions could not be replayed; see {ledger_file.relative_to(ROOT)}.remainder.jsonl")
    with out.open("rb") as fh:
        applied = sh(f"{a.app_exec} staging-refresh apply --from -", stdin=fh, cwd=app_cwd, check=False)
    if applied.returncode == 2:
        fail(f"apply refused the export; the rows are safe in {out.relative_to(ROOT)}", 2)
    if applied.returncode == 1:
        print(f"some rows could not be put back; they are still in {out.relative_to(ROOT)} (see the report)")

    # 5. production's sessions are not staging's
    if not a.keep_sessions:
        sh(f"{a.app_exec} sessions revoke --all", cwd=app_cwd)
    print(f"done; the export is kept at {out.relative_to(ROOT)}")
    return applied.returncode


if __name__ == "__main__":
    sys.exit(main())