#!/usr/bin/env python3
"""Provision ZITADEL end to end: machine key, Terraform, SOPS secrets, users.

One command replaces the manual sequence the Authentik script
(scripts/iam/provision-rbac.py) used to do in a single pass. Phases run in a
fixed order and every phase is safe to re-run:

  keys         extract the iam-admin machine key and PAT from namespace iam
  apply        terraform apply terraform/zitadel (project, roles, apps, machines)
  secrets      write generated client ids/secrets back into the SOPS files
  users        import one human user, converting the Django password hash
  verify       check discovery, client count and the groups-claim gate

`all` runs keys, apply, secrets, verify. It never runs users: the password hash
must be supplied by a human, one user at a time.

Nothing here talks to Authentik. Authentik keeps serving until a consumer is
moved deliberately.

Usage:
  scripts/iam/zitadel-provision.py keys
  scripts/iam/zitadel-provision.py apply --org-id <id>
  scripts/iam/zitadel-provision.py secrets --dry-run
  scripts/iam/zitadel-provision.py secrets --apply-config
  scripts/iam/zitadel-provision.py users --django-hash '<pbkdf2_sha256$...>'
  scripts/iam/zitadel-provision.py verify
"""

from __future__ import annotations

import argparse
import base64
import json
import os
import re
import signal
import stat
import subprocess
import sys
import tempfile
from pathlib import Path

import yaml

REPO = Path(__file__).resolve().parents[2]
TF_DIR = REPO / "terraform" / "zitadel"
SOPS_DIR = REPO / "k8s" / "argocd" / "secrets"
SOPS_CONFIG = REPO / ".sops.yaml"
CONF_DIR = Path(os.environ.get("ZITADEL_CONF_DIR", Path.home() / ".config" / "zitadel"))
NAMESPACE = os.environ.get("ZITADEL_NAMESPACE", "iam")
BASE_URL = os.environ.get("ZITADEL_BASE_URL", "https://zitadel.riotpiao.com")

_TEMP_FILES: set[str] = set()


def _cleanup_temps(*_a) -> None:
    for path in list(_TEMP_FILES):
        try:
            os.unlink(path)
        except OSError:
            pass
        _TEMP_FILES.discard(path)


for _sig in (signal.SIGTERM, signal.SIGINT, signal.SIGHUP):
    signal.signal(_sig, lambda s, f: (_cleanup_temps(), sys.exit(130)))


def die(msg: str) -> None:
    _cleanup_temps()
    print(f"FATAL: {msg}", file=sys.stderr)
    sys.exit(1)


def run(cmd: list[str], **kw) -> subprocess.CompletedProcess:
    return subprocess.run(cmd, check=False, text=True, capture_output=True, **kw)


def need(binary: str) -> None:
    if subprocess.run(["sh", "-c", f"command -v {binary}"], capture_output=True).returncode != 0:
        die(f"{binary} is not on PATH")


def write_private(path: Path, data: bytes) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    fd = os.open(path, os.O_WRONLY | os.O_CREAT | os.O_TRUNC, 0o600)
    try:
        os.fchmod(fd, 0o600)  # open() mode applies at creation only
        with os.fdopen(fd, "wb") as fh:
            fh.write(data)
    finally:
        pass


def secret_value(name: str, key: str) -> str:
    cp = run(["kubectl", "-n", NAMESPACE, "get", "secret", name, "-o", f"jsonpath={{{key}}}"])
    if cp.returncode != 0:
        die(f"cannot read secret {NAMESPACE}/{name}: {cp.stderr.strip()}")
    raw = cp.stdout.strip()
    if not raw:
        die(f"secret {NAMESPACE}/{name} has no key {key}")
    return base64.b64decode(raw).decode()


def curl_json(method: str, url: str, token: str | None = None, body: dict | None = None,
              timeout: int = 20) -> tuple[int, str]:
    """Call the API without putting a secret in the process argv.

    `curl --config -` reads the options from stdin, so neither the bearer token
    nor the password hash of a user import shows up in `ps`.
    """
    def esc(value: str) -> str:
        return value.replace("\\", "\\\\").replace('"', '\\"')

    cfg = [f'url = "{esc(url)}"', f'request = "{method}"', "silent", "show-error",
           f"max-time = {timeout}", 'write-out = "\\n%{http_code}"',
           'header = "Content-Type: application/json"']
    if token:
        cfg.append(f'header = "Authorization: Bearer {esc(token)}"')
    if body is not None:
        cfg.append(f'data = "{esc(json.dumps(body))}"')
    cp = run(["curl", "--config", "-"], input="\n".join(cfg) + "\n")
    if cp.returncode != 0:
        die(f"{method} {url} failed: {cp.stderr.strip()}")
    payload, _, code = cp.stdout.rpartition("\n")
    return int(code.strip() or 0), payload


# ── phase: keys ────────────────────────────────────────────────────────────

def phase_keys(_args: argparse.Namespace) -> int:
    write_private(CONF_DIR / "iam-admin.json", secret_value("iam-admin", ".data.iam-admin\\.json").encode())
    write_private(CONF_DIR / "iam-admin.pat", secret_value("iam-admin-pat", ".data.pat").strip().encode())
    print(f"machine key  -> {CONF_DIR / 'iam-admin.json'} (0600)")
    print(f"PAT          -> {CONF_DIR / 'iam-admin.pat'} (0600)")
    print("both files stay outside the repo; never commit them")
    return 0


# ── phase: apply ───────────────────────────────────────────────────────────

def phase_apply(args: argparse.Namespace) -> int:
    need("terraform")
    if not args.org_id:
        die("apply requires --org-id (ZITADEL console -> Organization settings -> Resource ID)")
    key = CONF_DIR / "iam-admin.json"
    if not key.exists():
        die(f"{key} missing — run the keys phase first")
    env = dict(os.environ, ZITADEL_JWT_PROFILE_FILE=str(key))

    for cmd in (["init", "-input=false"], ["fmt", "-check", "-recursive"]):
        cp = run(["terraform", f"-chdir={TF_DIR}", *cmd], env=env)
        if cp.returncode != 0:
            die(f"terraform {' '.join(cmd)} failed: {cp.stderr.strip() or cp.stdout.strip()}")

    cp = run(["terraform", f"-chdir={TF_DIR}", "apply", "-input=false", "-auto-approve",
              f"-var=org_id={args.org_id}"], env=env)
    sys.stdout.write(cp.stdout)
    sys.stderr.write(cp.stderr)
    if cp.returncode != 0:
        die("terraform apply failed — no SOPS value was written. Objects created before the "
            "failure stay in ZITADEL; re-run apply once the cause is fixed")
    return 0


# ── phase: secrets ─────────────────────────────────────────────────────────

# client -> [(sops file, Secret name inside the file, key, kind)]
# Only secrets with a git-tracked SOPS file appear here. Every other client is
# printed as manual: guessing a path writes a credential into the wrong Secret.
SECRET_MAP = {
    "grafana": [("grafana-oidc-secrets.enc.yaml", "grafana-oidc",
                 "GF_AUTH_GENERIC_OAUTH_CLIENT_SECRET", "secret")],
    "minio": [("minio-secrets.enc.yaml", "minio-oidc",
               "MINIO_IDENTITY_OPENID_CLIENT_SECRET", "secret")],
    "paperless-ai-agent": [("paperless-ai-secrets.enc.yaml", "paperless-ai-config",
                            "LLM_AUTH_CLIENT_SECRET", "secret")],
}

# Clients whose client_id lives in a plain (unencrypted) manifest. They need the
# id update as much as the secret: ZITADEL generates a new client_id for every
# application, so a secret-only re-point yields invalid_client.
CLIENT_ID_MAP = {
    "grafana": ("k8s/infra/logging/grafana-values.yaml",
                r"(?m)^(\s*client_id:\s*)\S+\s*$", r"\g<1>{value}"),
    "minio": ("k8s/infra/minio/minio-tenant.yaml",
              r"(?m)^(\s*value:\s*)\"minio\"\s*$", r'\g<1>"{value}"'),
    "paperless-ai-agent": ("k8s/apps/paperless/paperless-ai.yaml",
                           r"client_id=paperless-ai-agent", "client_id={value}"),
}

# Secrets that a live workload mounts but no SOPS file in git defines. The phase
# cannot re-point them safely, so it names them.
MANUAL_ONLY = {
    "homarr": "Secret homarr-oidc (keys client-id, client-secret) has no SOPS file in git — "
              "adopt it first, then add it to SECRET_MAP. The repo also carries an orphan "
              "homarr-auth-oidc.enc.yaml that nothing consumes.",
}


def terraform_outputs() -> dict:
    need("terraform")
    key = CONF_DIR / "iam-admin.json"
    if not key.exists():
        die(f"{key} missing — run the keys phase first")
    env = dict(os.environ, ZITADEL_JWT_PROFILE_FILE=str(key))
    cp = run(["terraform", f"-chdir={TF_DIR}", "output", "-json"], env=env)
    if cp.returncode != 0:
        die(f"terraform output failed: {cp.stderr.strip()}")
    return json.loads(cp.stdout or "{}")


def sops_env() -> dict:
    env = dict(os.environ)
    env.setdefault("SOPS_AGE_KEY_FILE", str(Path.home() / ".sops" / "key.txt"))
    return env


def write_secret(rel_path: str, secret_name: str, key: str, value: str,
                 dry_run: bool) -> None:
    path = SOPS_DIR / rel_path
    if not path.exists():
        die(f"{path} does not exist — the mapping and the repo disagree")
    env = sops_env()
    dec = run(["sops", "--config", str(SOPS_CONFIG), "decrypt", str(path)], env=env)
    if dec.returncode != 0:
        die(f"cannot decrypt {rel_path}: {dec.stderr.strip()}")
    docs = [d for d in yaml.safe_load_all(dec.stdout) if d]
    matching = [d for d in docs if (d.get("metadata") or {}).get("name") == secret_name]
    if len(matching) != 1:
        die(f"{rel_path}: expected exactly one document named {secret_name}, found {len(matching)}")

    current = (matching[0].get("stringData") or {}).get(key)
    if current == value:
        print(f"  unchanged {rel_path}:{secret_name}:{key}")
        return
    if dry_run:
        print(f"  would set {rel_path}:{secret_name}:{key} (len {len(value)})")
        return

    if len(docs) == 1:
        cp = run(["sops", "--config", str(SOPS_CONFIG), "set", "--value-stdin", "--idempotent",
                  str(path), f'["stringData"]["{key}"]'], env=env, input=json.dumps(value))
        if cp.returncode != 0:
            die(f"sops set failed on {rel_path}: {cp.stderr.strip()}")
        print(f"  wrote {rel_path}:{secret_name}:{key}")
        return

    # Multi-document file. `sops set` cannot address one document, so the whole
    # file is rewritten. sops encrypt reads the creation rule from the FILENAME,
    # so the temp path is overridden with the repo-relative path; without that
    # override sops finds no matching rule and refuses.
    matching[0].setdefault("stringData", {})[key] = value
    plain = "---\n".join(yaml.safe_dump(d, sort_keys=False) for d in docs)
    fd, tmp = tempfile.mkstemp(suffix=".yaml", dir=tempfile.gettempdir())
    _TEMP_FILES.add(tmp)
    try:
        with os.fdopen(fd, "w") as fh:
            fh.write(plain)
        enc = run(["sops", "--config", str(SOPS_CONFIG), "encrypt",
                   "--filename-override", str(path), tmp], env=env)
        if enc.returncode != 0:
            die(f"sops encrypt failed for {rel_path}: {enc.stderr.strip()}")
        mode = stat.S_IMODE(path.stat().st_mode)
        fd2, staged = tempfile.mkstemp(suffix=".enc.yaml", dir=str(path.parent))
        try:
            with os.fdopen(fd2, "w") as fh:
                fh.write(enc.stdout)
            os.chmod(staged, mode)
            os.replace(staged, path)  # same directory: atomic, never EXDEV
        finally:
            if os.path.exists(staged):
                os.unlink(staged)
        print(f"  wrote {rel_path}:{secret_name}:{key} (multi-document rewrite)")
    finally:
        if os.path.exists(tmp):
            os.unlink(tmp)
        _TEMP_FILES.discard(tmp)


def update_client_id(client: str, value: str, dry_run: bool, apply_config: bool) -> None:
    spec = CLIENT_ID_MAP.get(client)
    if not spec:
        return
    rel_path, pattern, repl = spec
    path = REPO / rel_path
    if not path.exists():
        die(f"{path} does not exist — the client-id mapping is stale")
    text = path.read_text()
    matches = list(re.finditer(pattern, text))
    if len(matches) != 1:
        die(f"{rel_path}: expected 1 match for the {client} client id, found {len(matches)}")
    new_line = matches[0].expand(repl.format(value=value))
    if matches[0].group(0) == new_line:
        print(f"  unchanged {rel_path}: {client} client id")
        return
    if dry_run or not apply_config:
        print(f"  would change {rel_path}: {matches[0].group(0).strip()} -> {new_line.strip()}")
        if not apply_config:
            print("    re-run with --apply-config to write this plaintext manifest edit")
        return
    path.write_text(text[:matches[0].start()] + new_line + text[matches[0].end():])
    print(f"  wrote {rel_path}: {client} client id")


def phase_secrets(args: argparse.Namespace) -> int:
    need("sops")
    out = terraform_outputs()
    ids = out.get("client_ids", {}).get("value", {})
    secrets = out.get("client_secrets", {}).get("value", {})
    machine_keys = out.get("machine_keys", {}).get("value", {})

    if not ids and not machine_keys:
        die("terraform has no ZITADEL outputs yet — run the apply phase first")

    print("[1/3] client secrets into SOPS")
    for client, entries in SECRET_MAP.items():
        for rel_path, secret_name, key, kind in entries:
            value = ids.get(client) if kind == "id" else secrets.get(client)
            if value:
                write_secret(rel_path, secret_name, key, value, args.dry_run)

    print("[2/3] client ids in plain manifests")
    for client in list(CLIENT_ID_MAP) + list(SECRET_MAP):
        if ids.get(client):
            update_client_id(client, ids[client], args.dry_run, args.apply_config)

    print("[3/3] machine keys")
    key_dir = CONF_DIR / "machine-keys"
    for name, key_json in machine_keys.items():
        if args.dry_run:
            print(f"  would write {key_dir / (name + '.json')}")
        else:
            write_private(key_dir / f"{name}.json", key_json.encode())
            print(f"  wrote {key_dir / (name + '.json')} (0600)")

    unhandled = sorted(set(ids) - set(SECRET_MAP) - set(CLIENT_ID_MAP) - set(MANUAL_ONLY))
    if MANUAL_ONLY or unhandled:
        print("\n  MANUAL — no safe write path exists for these clients:")
        for client, why in MANUAL_ONLY.items():
            print(f"    {client}: {why}")
        for client in unhandled:
            print(f"    {client}: set its client_id and secret in its own manifest or Secret file")
    return 0


# ── phase: users ───────────────────────────────────────────────────────────

HASH_RE = re.compile(r"^pbkdf2_sha256\$(?P<rounds>\d+)\$(?P<salt>[^$]+)\$(?P<hash>[^$]+)$")


def convert_hash(django_hash: str) -> str:
    """Django pbkdf2_sha256 -> passwap MCF $pbkdf2-sha256$...

    ZITADEL imports one MCF string via hashedPassword.hash (max 200 chars) and
    pins passwap, whose pbkdf2 format is $pbkdf2-sha256$<rounds>$<salt>$<hash>
    with alt-base64 (standard base64, + -> ., no padding). Rounds, salt and hash
    bytes are unchanged, so the password still verifies.
    """
    m = HASH_RE.match(django_hash.strip())
    if not m:
        die("unsupported hash format — only pbkdf2_sha256 is convertible")
    alt = m.group("hash").rstrip("=").replace("+", ".")
    mcf = f"$pbkdf2-sha256${m.group('rounds')}${m.group('salt')}${alt}"
    if len(mcf) > 200:
        die(f"converted hash is {len(mcf)} chars, ZITADEL accepts 200")
    return mcf


def phase_users(args: argparse.Namespace) -> int:
    pat_file = CONF_DIR / "iam-admin.pat"
    if not pat_file.exists():
        die(f"{pat_file} missing — run the keys phase first")
    if not args.django_hash:
        die("pass --django-hash with the pbkdf2_sha256 value exported from Authentik")
    body = {
        "username": args.username,
        "profile": {"givenName": args.given_name, "familyName": args.family_name},
        "email": {"email": args.email, "isVerified": True},
        "hashedPassword": {"hash": convert_hash(args.django_hash)},
        "passwordChangeRequired": False,
    }
    code, payload = curl_json("POST", f"{BASE_URL}/v2/users/human",
                              token=pat_file.read_text().strip(), body=body)
    if code in (200, 201):
        print(f"  {args.username}: created")
        return 0
    if code == 409:
        print(f"  {args.username}: already exists — the password was NOT updated")
        return 1
    die(f"create {args.username} -> {code} {payload[:300]}")


# ── phase: verify ──────────────────────────────────────────────────────────

def phase_verify(_args: argparse.Namespace) -> int:
    code, _ = curl_json("GET", f"{BASE_URL}/.well-known/openid-configuration", timeout=10)
    print(f"discovery      : {code}")
    if code != 200:
        die("discovery endpoint is not 200 — the instance or the ingress is not ready")

    out = terraform_outputs()
    ids = out.get("client_ids", {}).get("value", {})
    print(f"client apps    : {len(ids)}")
    enabled = bool(out.get("groups_claim_enabled", {}).get("value"))
    print(f"groups action  : {enabled}")
    if not enabled:
        print("BLOCKED for cutover: the groups claim action is disabled. Every consumer")
        print("that reads `groups` (kube-apiserver, MinIO, Grafana) sees no groups.")
        print("Enable it only after the live gate passes, with the script path set.")
    return 0


PHASES = {"keys": phase_keys, "apply": phase_apply, "secrets": phase_secrets,
          "users": phase_users, "verify": phase_verify}


def main() -> int:
    parser = argparse.ArgumentParser(description="Provision ZITADEL end to end.")
    parser.add_argument("phase", choices=[*PHASES, "all"])
    parser.add_argument("--org-id", help="ZITADEL org id (required for apply)")
    parser.add_argument("--dry-run", action="store_true", help="print actions, write nothing")
    parser.add_argument("--apply-config", action="store_true",
                        help="also rewrite client ids inside plain manifests")
    parser.add_argument("--username", default="rock")
    parser.add_argument("--email", default="rock@riotpiao.com")
    parser.add_argument("--given-name", default="Rock")
    parser.add_argument("--family-name", default="Liang")
    parser.add_argument("--django-hash", help="pbkdf2_sha256 value exported from Authentik")
    args = parser.parse_args()

    need("kubectl")
    if args.phase == "all":
        if not args.org_id:
            die("all requires --org-id")
        for name in ("keys", "apply"):
            rc = PHASES[name](args)
            if rc:
                return rc
        rc = phase_secrets(args)
        if rc:
            return rc
        return phase_verify(args)
    return PHASES[args.phase](args)


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