scripts/admin_deploy.py

"""Admin / safe-deploy CLI for the live Kink Discovery Space.

Every operation a user can do (and several they can't) — exposed as subcommands. Used to:

  * back up every profile to a local file before a risky deploy
  * restore those profiles after a deploy if anything goes sideways
  * fix one user without redeploying (lost token, stuck partner state, bad rating)
  * inspect the live state at a glance (stats, health, user list)

Backed by ``/admin/*`` endpoints on the live API. Auth: ``KINK_ADMIN_SECRET`` env var (or
``--admin-secret``). The Space must have the same secret in its environment — bootstrap with
``admin_deploy.py provision-secret`` once.

Examples
~~~~~~~~

    # one-time: install KINK_ADMIN_SECRET as a Space secret
    .venv/bin/python scripts/admin_deploy.py provision-secret

    # day-to-day
    .venv/bin/python scripts/admin_deploy.py list-users
    .venv/bin/python scripts/admin_deploy.py user-info real-bunny-72
    .venv/bin/python scripts/admin_deploy.py reset-token real-bunny-72
    .venv/bin/python scripts/admin_deploy.py stats
    .venv/bin/python scripts/admin_deploy.py snapshot-push

    # safe deploy (backup → publish → wait → snapshot push → snapshot verify)
    .venv/bin/python scripts/admin_deploy.py safe-deploy

    # restore a known-good backup over whatever's live
    .venv/bin/python scripts/admin_deploy.py restore data/backups/2026-04-30T1730.db
"""
from __future__ import annotations

import argparse
import json
import os
import secrets
import sys
import time
import urllib.error
import urllib.parse
import urllib.request
from datetime import datetime, timezone
from pathlib import Path

REPO = Path(__file__).resolve().parent.parent
sys.path.insert(0, str(REPO))

from backend.repo_dotenv import load_repo_dotenv  # noqa: E402

DEFAULT_BASE = os.environ.get("KINK_AUDIT_BASE_URL", "https://perplexed7675-kink-discovery.hf.space")
DEFAULT_SPACE = os.environ.get("HF_SPACE_REPO", "Perplexed7675/kink-discovery")
BACKUP_DIR = REPO / "data" / "backups"


# ─────────────────────────────────────────────────────────────────────────────
# HTTP helpers (stdlib only — keeps the script self-contained on a fresh checkout)
# ─────────────────────────────────────────────────────────────────────────────


def _admin_headers(secret: str) -> dict[str, str]:
    return {"x-admin-secret": secret}


def _request(method: str, url: str, *, headers: dict | None = None, data: bytes | None = None,
             content_type: str | None = None, timeout: int = 120) -> tuple[int, bytes]:
    headers = dict(headers or {})
    if content_type:
        headers["Content-Type"] = content_type
    req = urllib.request.Request(url, data=data, method=method, headers=headers)
    try:
        with urllib.request.urlopen(req, timeout=timeout) as resp:
            return resp.status, resp.read()
    except urllib.error.HTTPError as e:
        return e.code, e.read()


def _api(method: str, base: str, path: str, secret: str, **kw) -> tuple[int, bytes]:
    return _request(method, base.rstrip("/") + path, headers=_admin_headers(secret), **kw)


def _json_or_raise(status: int, body: bytes, label: str) -> dict | list:
    if 200 <= status < 300:
        return json.loads(body) if body else {}
    raise SystemExit(f"{label}: HTTP {status} — {body.decode(errors='replace')[:400]}")


# ─────────────────────────────────────────────────────────────────────────────
# Subcommands
# ─────────────────────────────────────────────────────────────────────────────


def cmd_list_users(args) -> int:
    s, body = _api("GET", args.base_url, "/admin/users?limit=2000", args.admin_secret)
    data = _json_or_raise(s, body, "list-users")
    items = data["items"]
    print(f"{len(items)} users (showing all):")
    print(f"  {'user_id':<28} {'plays':>6} {'scen':>5} {'roles':>5} {'links':>5} {'groups':>6}")
    for it in items:
        print(f"  {it['user_id']:<28} {it['plays']:>6} {it['scenarios']:>5} {it['roles']:>5} "
              f"{it['partner_links']:>5} {it['group_memberships']:>6}")
    return 0


def cmd_user_info(args) -> int:
    s, body = _api("GET", args.base_url, f"/admin/users/{urllib.parse.quote(args.user_id)}", args.admin_secret)
    data = _json_or_raise(s, body, "user-info")
    print(json.dumps(data, indent=2, sort_keys=True))
    return 0


def cmd_reset_token(args) -> int:
    s, body = _api("POST", args.base_url,
                   f"/admin/users/{urllib.parse.quote(args.user_id)}/regenerate-token",
                   args.admin_secret)
    data = _json_or_raise(s, body, "reset-token")
    print(json.dumps(data, indent=2))
    return 0


def cmd_delete_user(args) -> int:
    if not args.yes:
        print(f"Refusing without --yes; will delete {args.user_id} permanently.", file=sys.stderr)
        return 2
    s, body = _api("DELETE", args.base_url, f"/admin/users/{urllib.parse.quote(args.user_id)}",
                   args.admin_secret)
    data = _json_or_raise(s, body, "delete-user")
    print(json.dumps(data, indent=2))
    return 0


def cmd_scrub_plays(args) -> int:
    if not args.yes:
        print(f"Refusing without --yes; will wipe ALL plays for {args.user_id}.", file=sys.stderr)
        return 2
    s, body = _api("POST", args.base_url,
                   f"/admin/users/{urllib.parse.quote(args.user_id)}/scrub-plays",
                   args.admin_secret)
    print(json.dumps(_json_or_raise(s, body, "scrub-plays"), indent=2))
    return 0


def cmd_delete_play(args) -> int:
    s, body = _api("DELETE", args.base_url,
                   f"/admin/users/{urllib.parse.quote(args.user_id)}/plays/{urllib.parse.quote(args.kink_id)}",
                   args.admin_secret)
    print(json.dumps(_json_or_raise(s, body, "delete-play"), indent=2))
    return 0


def cmd_force_link(args) -> int:
    s, body = _api("POST", args.base_url,
                   f"/admin/users/{urllib.parse.quote(args.user_id)}/partners/{urllib.parse.quote(args.partner_id)}/force-link",
                   args.admin_secret)
    print(json.dumps(_json_or_raise(s, body, "force-link"), indent=2))
    return 0


def cmd_unlink(args) -> int:
    s, body = _api("DELETE", args.base_url,
                   f"/admin/partners/{urllib.parse.quote(args.user_a)}/{urllib.parse.quote(args.user_b)}",
                   args.admin_secret)
    print(json.dumps(_json_or_raise(s, body, "unlink"), indent=2))
    return 0


def cmd_backup(args) -> int:
    BACKUP_DIR.mkdir(parents=True, exist_ok=True)
    # ``--out`` is only on the ``backup`` subparser; ``safe-deploy`` calls cmd_backup directly
    # with its own Namespace that has no ``out`` attribute. Use getattr to keep both flows working.
    out_arg = getattr(args, "out", None)
    if out_arg:
        out = Path(out_arg)
    else:
        ts = datetime.now(timezone.utc).strftime("%Y-%m-%dT%H%M%SZ")
        out = BACKUP_DIR / f"user_state_{ts}.db"
    s, body = _api("GET", args.base_url, "/admin/export", args.admin_secret, timeout=300)
    if not (200 <= s < 300):
        raise SystemExit(f"backup: HTTP {s} — {body[:400]!r}")
    out.write_bytes(body)
    print(f"wrote {out}  ({len(body):,} bytes)")
    return 0


def cmd_restore(args) -> int:
    if not args.yes:
        print(f"Refusing without --yes; will overwrite live user-state from {args.file!r}.", file=sys.stderr)
        return 2
    body = Path(args.file).read_bytes()
    s, resp = _api("POST", args.base_url, "/admin/import", args.admin_secret,
                   data=body, content_type="application/octet-stream", timeout=300)
    print(json.dumps(_json_or_raise(s, resp, "restore"), indent=2))
    return 0


def cmd_snapshot_push(args) -> int:
    s, body = _api("POST", args.base_url, "/admin/snapshot/push", args.admin_secret, timeout=180)
    print(json.dumps(_json_or_raise(s, body, "snapshot-push"), indent=2))
    return 0


def cmd_snapshot_pull(args) -> int:
    if not args.yes:
        print("Refusing without --yes; will overwrite local store with Hub snapshot.", file=sys.stderr)
        return 2
    s, body = _api("POST", args.base_url, "/admin/snapshot/pull", args.admin_secret, timeout=180)
    print(json.dumps(_json_or_raise(s, body, "snapshot-pull"), indent=2))
    return 0


def cmd_stats(args) -> int:
    s, body = _api("GET", args.base_url, "/admin/stats", args.admin_secret)
    print(json.dumps(_json_or_raise(s, body, "stats"), indent=2))
    return 0


def cmd_health(args) -> int:
    s, body = _api("GET", args.base_url, "/admin/health", args.admin_secret)
    print(json.dumps(_json_or_raise(s, body, "health"), indent=2))
    return 0


def cmd_provision_secret(args) -> int:
    """Generate a fresh KINK_ADMIN_SECRET, write it to .env.local AND set it as a Space secret."""
    load_repo_dotenv()
    hf_token = os.environ.get("HF_TOKEN") or os.environ.get("HUGGING_FACE_HUB_TOKEN")
    if not hf_token:
        raise SystemExit("provision-secret needs HF_TOKEN in env / .env.local")
    new_secret = args.admin_secret or os.environ.get("KINK_ADMIN_SECRET") or secrets.token_urlsafe(32)

    from huggingface_hub import HfApi
    api = HfApi(token=hf_token)
    api.add_space_secret(repo_id=DEFAULT_SPACE, key="KINK_ADMIN_SECRET", value=new_secret)
    print(f"set KINK_ADMIN_SECRET on {DEFAULT_SPACE} (Space will pick it up on next restart)")

    env_local = REPO / ".env.local"
    contents = env_local.read_text(encoding="utf-8") if env_local.is_file() else ""
    lines = [ln for ln in contents.splitlines() if not ln.startswith("KINK_ADMIN_SECRET=")]
    lines.append(f"KINK_ADMIN_SECRET={new_secret}")
    env_local.write_text("\n".join(lines) + "\n", encoding="utf-8")
    print(f"wrote KINK_ADMIN_SECRET to {env_local}")
    print("(restart the Space to pick up the new secret)")
    return 0


def cmd_safe_deploy(args) -> int:
    """backup → publish → wait for RUNNING → verify users restored → manual snapshot push.

    Idempotent on retries: the backup is timestamped so each call snapshots state at that
    moment, and the post-deploy verify just compares user IDs (no destructive ops).
    """
    print("== safe-deploy ==", flush=True)
    print("step 1/5: backup", flush=True)
    if cmd_backup(args) != 0:
        return 1

    print("\nstep 2/5: collect pre-deploy user count", flush=True)
    s, body = _api("GET", args.base_url, "/admin/users?limit=2000", args.admin_secret)
    pre_users = {u["user_id"] for u in _json_or_raise(s, body, "pre-count")["items"]}
    print(f"  pre-deploy users: {len(pre_users)}", flush=True)

    print("\nstep 3/5: publish to HF Space", flush=True)
    import subprocess
    rc = subprocess.run(
        [sys.executable, str(REPO / "scripts" / "publish_hf_space.py"), "--no-wait"],
        check=False,
    ).returncode
    if rc != 0:
        raise SystemExit(f"publish exited {rc}")

    print("\nstep 4/5: wait for the new build to come up RUNNING + boot-restore to log", flush=True)
    load_repo_dotenv()
    hf_token = os.environ.get("HF_TOKEN") or os.environ.get("HUGGING_FACE_HUB_TOKEN")
    if not hf_token:
        raise SystemExit("HF_TOKEN missing — cannot poll Space stage")
    deadline = time.time() + 600
    last = None
    while time.time() < deadline:
        req = urllib.request.Request(
            f"https://huggingface.co/api/spaces/{DEFAULT_SPACE}",
            headers={"Authorization": f"Bearer {hf_token}"},
        )
        runtime = json.loads(urllib.request.urlopen(req, timeout=30).read())["runtime"]
        stage = runtime["stage"]
        if stage != last:
            print(f"  stage={stage}", flush=True)
            last = stage
        if stage == "RUNNING":
            try:
                r = urllib.request.urlopen(args.base_url.rstrip('/') + "/health", timeout=20).read()
                if json.loads(r).get("user_snapshot", {}).get("enabled"):
                    break
            except Exception:
                pass
        time.sleep(8)
    else:
        raise SystemExit("timeout waiting for build")

    print("\nstep 5/5: verify users restored + push fresh snapshot", flush=True)
    s, body = _api("GET", args.base_url, "/admin/users?limit=2000", args.admin_secret)
    post_users = {u["user_id"] for u in _json_or_raise(s, body, "post-count")["items"]}
    print(f"  post-deploy users: {len(post_users)}", flush=True)
    missing = pre_users - post_users
    extra = post_users - pre_users
    if missing:
        print(f"  WARNING: {len(missing)} pre-deploy users are MISSING post-deploy: {sorted(missing)[:10]}",
              flush=True)
        print("  → run: admin_deploy.py restore <backup-file> --yes", flush=True)
    else:
        print("  ✓ all pre-deploy users present after deploy", flush=True)
    if extra:
        print(f"  (note: {len(extra)} new users showed up post-deploy; that's fine)", flush=True)

    print("\n  pushing fresh snapshot to Hub…", flush=True)
    s, body = _api("POST", args.base_url, "/admin/snapshot/push", args.admin_secret, timeout=180)
    print("  " + json.dumps(_json_or_raise(s, body, "snapshot-push"), indent=2).replace("\n", "\n  "))

    return 0 if not missing else 3


# ─────────────────────────────────────────────────────────────────────────────
# argparse wiring
# ─────────────────────────────────────────────────────────────────────────────


def main() -> int:
    load_repo_dotenv()
    p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    p.add_argument("--base-url", default=DEFAULT_BASE)
    p.add_argument("--admin-secret", default=os.environ.get("KINK_ADMIN_SECRET", ""),
                   help="Defaults to $KINK_ADMIN_SECRET")

    sub = p.add_subparsers(dest="cmd", required=True)

    sub.add_parser("list-users", help="List all profiles with summary counts").set_defaults(fn=cmd_list_users)

    s_info = sub.add_parser("user-info", help="Full state for one profile (incl. private_token)")
    s_info.add_argument("user_id")
    s_info.set_defaults(fn=cmd_user_info)

    s_reset = sub.add_parser("reset-token", help="Issue a new private_token for a user who lost theirs")
    s_reset.add_argument("user_id")
    s_reset.set_defaults(fn=cmd_reset_token)

    s_del = sub.add_parser("delete-user", help="Hard-delete a profile + all its state")
    s_del.add_argument("user_id")
    s_del.add_argument("--yes", action="store_true")
    s_del.set_defaults(fn=cmd_delete_user)

    s_scrub = sub.add_parser("scrub-plays", help="Wipe all PlayPreference for one user")
    s_scrub.add_argument("user_id")
    s_scrub.add_argument("--yes", action="store_true")
    s_scrub.set_defaults(fn=cmd_scrub_plays)

    s_dp = sub.add_parser("delete-play", help="Remove ONE play preference")
    s_dp.add_argument("user_id")
    s_dp.add_argument("kink_id")
    s_dp.set_defaults(fn=cmd_delete_play)

    s_link = sub.add_parser("force-link", help="Force-create a partner link without the request flow")
    s_link.add_argument("user_id")
    s_link.add_argument("partner_id")
    s_link.set_defaults(fn=cmd_force_link)

    s_unlink = sub.add_parser("unlink", help="Drop a partner link between two users")
    s_unlink.add_argument("user_a")
    s_unlink.add_argument("user_b")
    s_unlink.set_defaults(fn=cmd_unlink)

    s_b = sub.add_parser("backup", help="Download a user_state.db snapshot from live")
    s_b.add_argument("--out", help="Path; default data/backups/user_state_<ts>.db")
    s_b.set_defaults(fn=cmd_backup)

    s_r = sub.add_parser("restore", help="Upload a user_state.db and INSERT OR REPLACE")
    s_r.add_argument("file")
    s_r.add_argument("--yes", action="store_true")
    s_r.set_defaults(fn=cmd_restore)

    sub.add_parser("snapshot-push", help="Force flush user-state to the Hub dataset").set_defaults(fn=cmd_snapshot_push)

    s_pull = sub.add_parser("snapshot-pull", help="Force pull Hub snapshot + INSERT OR REPLACE into live store")
    s_pull.add_argument("--yes", action="store_true")
    s_pull.set_defaults(fn=cmd_snapshot_pull)

    sub.add_parser("stats", help="Aggregate row counts + top kinks").set_defaults(fn=cmd_stats)
    sub.add_parser("health", help="Detailed admin health (incl. last_restore log)").set_defaults(fn=cmd_health)

    sub.add_parser(
        "provision-secret",
        help="Generate KINK_ADMIN_SECRET, write to .env.local, set as Space secret",
    ).set_defaults(fn=cmd_provision_secret)

    sub.add_parser(
        "safe-deploy",
        help="Backup → publish → wait → verify → push fresh snapshot",
    ).set_defaults(fn=cmd_safe_deploy)

    args = p.parse_args()
    if args.cmd != "provision-secret" and not args.admin_secret:
        raise SystemExit("KINK_ADMIN_SECRET missing — set it via env or --admin-secret, or run `provision-secret`.")
    return args.fn(args)


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