#!/usr/bin/env python3
"""Rotate the OIDC client secret of every owned application, one at a time.

This is the manual execution of what a KeyRotation object declares. It exists so
the credentials can rotate today, before the controller image is built, and so the
controller has a reference implementation to match.

For each application: generate a new secret in ZITADEL, write it into the
Kubernetes Secret, roll the consumer, then prove the consumer still talks to
ZITADEL. A failure stops the run, because the next application should not start
while one is in an unknown state.
"""
from __future__ import annotations

import json
import subprocess
import sys
import time
from pathlib import Path

sys.path.insert(0, str(Path(__file__).parent))
import importlib.util

spec = importlib.util.spec_from_file_location("zp", Path(__file__).with_name("zitadel-provision.py"))
zp = importlib.util.module_from_spec(spec)
spec.loader.exec_module(zp)

PROJECT = "396757796898471937"
ISSUER = "https://zitadel.riotpiao.com"

# app key, ZITADEL app id, Secret name, namespace, Secret key, consumer kind, consumer name
APPS = [
    ("grafana",   "grafana",   "grafana-oidc",         "logging",   "GF_AUTH_GENERIC_OAUTH_CLIENT_SECRET", "deploy", "grafana"),
    ("argocd",    "argocd",    "oidc-secret",           "argocd",    "client-secret",                       "deploy", "argocd-server"),
    ("comfyui",   "comfyui",   "comfyui-oauth2-proxy", "comfyui",   "client-secret",                       "deploy", "comfyui-oauth2-proxy"),
    ("minio",     "minio",     "minio-oidc",           "storage",   "MINIO_IDENTITY_OPENID_CLIENT_SECRET",  "sts",    "minio-cluster-az-a"),
    ("paperless", "paperless", "paperless-oidc",       "paperless", "client-secret",                       "deploy", "paperless"),
    ("immich",    "immich",    "immich-oidc",          "immich",    "client-secret",                       "deploy", "immich-server"),
    ("forgejo",   "forgejo",   "forgejo-oidc",         "cicd",      "CLIENT_SECRET",                       "deploy", "forgejo-gitea"),
]


def sh(*args: str, timeout: int = 120) -> tuple[int, str]:
    p = subprocess.run(args, capture_output=True, text=True, timeout=timeout)
    return p.returncode, (p.stdout + p.stderr).strip()


def state() -> dict:
    return json.load(open(Path.home() / ".config/zitadel/state.json"))


def app_ids(token: str) -> dict[str, str]:
    """Map clientId to the app resource id. The two are different, and the secret
    endpoint takes the resource id."""
    code, body = zp._api("POST", f"/management/v1/projects/{PROJECT}/apps/_search", token,
                         {"query": {"limit": 100, "offset": 0}})
    out = {}
    for app in body.get("result", []) or []:
        cid = (app.get("oidcConfig") or {}).get("clientId")
        if cid and app.get("id"):
            out[cid] = app["id"]
    return out


def new_secret(token: str, app_id: str) -> str | None:
    code, body = zp._api("POST", f"/management/v1/projects/{PROJECT}/apps/{app_id}/oidc_config/_generate_client_secret", token)
    if code != 200:
        print(f"    generate failed: {code} {str(body)[:120]}")
        return None
    return body.get("clientSecret")


def patch_secret(ns: str, name: str, key: str, value: str) -> bool:
    import base64
    enc = base64.b64encode(value.encode()).decode()
    patch = json.dumps({"data": {key: enc}})
    code, out = sh("kubectl", "-n", ns, "patch", "secret", name, "--type=merge", "-p", patch)
    if code != 0:
        print(f"    secret patch failed: {out[:120]}")
    return code == 0


def roll(kind: str, ns: str, name: str) -> bool:
    code, out = sh("kubectl", "-n", ns, "rollout", "restart", f"{kind}/{name}", timeout=120)
    if code != 0:
        print(f"    rollout restart failed: {out[:120]}")
        return False
    code, out = sh("kubectl", "-n", ns, "rollout", "status", f"{kind}/{name}", "--timeout=90s", timeout=120)
    return code == 0


def main() -> int:
    only = sys.argv[1:] or None
    token = zp.bootstrap_token()
    if not token:
        print("no bootstrap token")
        return 2
    cids = state()["client_ids"]["value"]
    ids = app_ids(token)
    print(f"resolved {len(ids)} app resource ids")
    done, failed = [], []
    for key, app_id, sname, ns, skey, kind, cname in APPS:
        if only and key not in only:
            continue
        print(f"[{key}] issuer={ISSUER} client_id={cids[key]}")
        rid = ids.get(cids[key])
        if not rid:
            print(f"[{key}] no app resource id for client {cids[key]}")
            failed.append(key)
            continue
        secret = new_secret(token, rid)
        if not secret:
            failed.append(key)
            continue
        if not patch_secret(ns, sname, skey, secret):
            failed.append(key)
            continue
        if not roll(kind, ns, cname):
            failed.append(key)
            print(f"[{key}] FAILED: consumer did not become ready")
            continue
        code, out = sh("kubectl", "-n", ns, "get", "pods", "-l",
                       f"app={cname}" if kind == "deploy" else f"app.kubernetes.io/name={cname}",
                       "--no-headers")
        ready = "Running" in out
        print(f"[{key}] rotated, consumer {'ready' if ready else 'NOT ready'}")
        (done if ready else failed).append(key)
    print(f"\nrotated: {len(done)} {done}\nfailed: {len(failed)} {failed}")
    return 0 if not failed else 1


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