#!/usr/bin/env bash
set -euo pipefail

OPENBAO_NAMESPACE="${OPENBAO_NAMESPACE:-openbao}"
OPENBAO_RELEASE="${OPENBAO_RELEASE:-openbao}"
KUBECTL="${KUBECTL:-kubectl}"
TOKEN_FILE="${OPENBAO_TOKEN_FILE:-}"
SSH_MOUNT="${OPENBAO_SSH_MOUNT:-ssh}"
ROLES_SPEC="${ROLES_SPEC:-}"
POLICY_DIR="${POLICY_DIR:-}"
CA_PUBKEY_OUT="${OPENBAO_SSH_CA_PUBKEY_OUT:-}"
K8S_CA_SECRET="${OPENBAO_SSH_CA_K8S_SECRET:-openbao-ssh-ca-pub}"
DRY_RUN=0

REPO_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
ROLES_SPEC="${ROLES_SPEC:-$REPO_DIR/openbao/ssh/roles-spec.yaml}"
POLICY_DIR="${POLICY_DIR:-$REPO_DIR/openbao/policies}"

usage() {
  cat <<'USAGE'
Usage: scripts/openbao-apply-ssh-engine.sh [--dry-run]

Applies the OpenBao SSH secrets engine for ops-warden signing:
  - enable ssh/ mount (idempotent)
  - write adm/agt/atm roles from openbao/ssh/roles-spec.yaml
  - load warden-sign policy
  - export CA public key (optional K8s secret + local file)

Requires initialized, unsealed OpenBao and a token with ssh engine admin
(typically platform-admin). Token from OPENBAO_TOKEN_FILE or prompt.

Env:
  OPENBAO_SSH_CA_PUBKEY_OUT  Write CA pubkey to this path (non-secret)
  OPENBAO_SSH_CA_K8S_SECRET  K8s secret name in openbao namespace (default: openbao-ssh-ca-pub)
USAGE
}

while [ "$#" -gt 0 ]; do
  case "$1" in
    --dry-run)
      DRY_RUN=1
      shift
      ;;
    -h|--help)
      usage
      exit 0
      ;;
    *)
      echo "ERROR: unknown argument: $1" >&2
      usage >&2
      exit 2
      ;;
  esac
done

pod="${OPENBAO_RELEASE}-0"
WARNINGS=0

warn() {
  WARNINGS=$((WARNINGS + 1))
  printf 'WARN: %s\n' "$*" >&2
}

read_token() {
  if [ "$DRY_RUN" -eq 1 ]; then
    printf 'dry-run-token\n'
    return
  fi
  if [ -n "$TOKEN_FILE" ]; then
    if [ ! -f "$TOKEN_FILE" ]; then
      echo "ERROR: OPENBAO_TOKEN_FILE does not exist: $TOKEN_FILE" >&2
      exit 1
    fi
    head -n 1 "$TOKEN_FILE"
    return
  fi
  local token
  read -r -s -p "OpenBao token: " token
  printf '\n' >&2
  printf '%s\n' "$token"
}

kubectl_exec() {
  # shellcheck disable=SC2086
  $KUBECTL "$@"
}

remote_bao() {
  local token="$1"
  shift
  if [ "$DRY_RUN" -eq 1 ]; then
    printf 'DRY-RUN: bao %s\n' "$*"
    return 0
  fi
  printf '%s\n' "$token" | kubectl_exec exec -i -n "$OPENBAO_NAMESPACE" "$pod" -- \
    sh -c 'read -r BAO_TOKEN; export BAO_TOKEN; exec bao "$@"' sh "$@"
}

write_policy() {
  local token="$1"
  local name="$2"
  local file="$3"
  if [ ! -f "$file" ]; then
    echo "ERROR: missing policy file: $file" >&2
    exit 1
  fi
  if [ "$DRY_RUN" -eq 1 ]; then
    printf 'DRY-RUN: bao policy write %s %s\n' "$name" "$file"
    return 0
  fi
  { printf '%s\n' "$token"; cat "$file"; } | kubectl_exec exec -i -n "$OPENBAO_NAMESPACE" "$pod" -- \
    sh -c 'read -r BAO_TOKEN; export BAO_TOKEN; bao policy write "$1" -' sh "$name"
}

ensure_default_issuer() {
  local token="$1"
  local issuers_out
  if [ "$DRY_RUN" -eq 1 ]; then
    printf 'DRY-RUN: bao read %s/config/issuers\n' "$SSH_MOUNT"
    printf 'DRY-RUN: bao write %s/config/ca generate_signing_key=true key_type=ed25519\n' "$SSH_MOUNT"
    return 0
  fi
  if issuers_out="$(remote_bao "$token" read "${SSH_MOUNT}/config/issuers" 2>&1)"; then
    printf 'OK: default SSH issuer already configured.\n'
    return 0
  fi
  case "$issuers_out" in
    *"no default issuer"*)
      remote_bao "$token" write "${SSH_MOUNT}/config/ca" \
        generate_signing_key=true key_type=ed25519
      printf 'OK: generated default SSH CA issuer.\n'
      ;;
    *)
      printf '%s\n' "$issuers_out" >&2
      echo "ERROR: failed to read SSH issuer configuration" >&2
      exit 1
      ;;
  esac
}

enable_ssh_engine() {
  local token="$1"
  local output status
  if output="$(remote_bao "$token" secrets enable -path="$SSH_MOUNT" ssh 2>&1)"; then
    printf '%s\n' "$output"
    return 0
  fi
  status=$?
  case "$output" in
    *"path is already in use"*)
      printf 'OK: %s/ SSH secrets engine is already enabled.\n' "$SSH_MOUNT"
      return 0
      ;;
    *)
      printf '%s\n' "$output" >&2
      echo "ERROR: failed to enable SSH engine (exit $status)" >&2
      exit 1
      ;;
  esac
}

apply_roles() {
  local token="$1"
  if [ ! -f "$ROLES_SPEC" ]; then
    echo "ERROR: roles spec not found: $ROLES_SPEC" >&2
    exit 1
  fi
  if [ "$DRY_RUN" -eq 1 ]; then
    python3 - "$ROLES_SPEC" "$SSH_MOUNT" <<'PY'
import sys, yaml
spec = yaml.safe_load(open(sys.argv[1]))
mount = sys.argv[2]
for name, params in (spec.get("roles") or {}).items():
    args = " ".join(f"{k}={v}" for k, v in params.items())
    print(f"DRY-RUN: bao write {mount}/roles/{name} {args}")
PY
    return 0
  fi
  python3 - "$token" "$ROLES_SPEC" "$SSH_MOUNT" "$OPENBAO_NAMESPACE" "$pod" "$KUBECTL" <<'PY'
import shlex
import subprocess
import sys

token, spec_path, mount, namespace, pod, kubectl = sys.argv[1:7]
import yaml

kubectl_parts = shlex.split(kubectl) if kubectl else ["kubectl"]
spec = yaml.safe_load(open(spec_path))
roles = spec.get("roles") or {}
for role_name, params in roles.items():
    args = [f"{k}={v}" for k, v in params.items()]
    cmd = ["bao", "write", f"{mount}/roles/{role_name}"] + args
    proc = subprocess.run(
        kubectl_parts + ["exec", "-i", "-n", namespace, pod, "--", "sh", "-c",
         "read -r BAO_TOKEN; export BAO_TOKEN; exec bao \"$@\"", "sh"] + cmd,
        input=(token + "\n").encode(),
        capture_output=True,
    )
    if proc.returncode != 0:
        sys.stderr.write(proc.stderr.decode())
        raise SystemExit(f"failed to write role {role_name}")
    print(f"OK: {mount}/roles/{role_name}")
PY
}

export_ca_pubkey() {
  local token="$1"
  if [ "$DRY_RUN" -eq 1 ]; then
    printf 'DRY-RUN: bao read -field=public_key %s/public_key\n' "$SSH_MOUNT"
    return 0
  fi
  local pubkey
  pubkey="$(remote_bao "$token" read -field=public_key "${SSH_MOUNT}/config/ca" 2>/dev/null || true)"
  if [ -z "$pubkey" ]; then
    pubkey="$(remote_bao "$token" read -field=public_key "${SSH_MOUNT}/public_key" 2>/dev/null || true)"
  fi
  if [ -z "$pubkey" ]; then
    warn "Could not read SSH CA public key from ${SSH_MOUNT}/config/ca or ${SSH_MOUNT}/public_key"
    return 0
  fi
  local fingerprint
  fingerprint="$(printf '%s' "$pubkey" | sha256sum | awk '{print $1}')"
  printf 'OK: SSH CA public key fingerprint sha256:%s\n' "$fingerprint"

  if [ -n "$CA_PUBKEY_OUT" ]; then
    mkdir -p "$(dirname "$CA_PUBKEY_OUT")"
    printf '%s\n' "$pubkey" >"$CA_PUBKEY_OUT"
    chmod 644 "$CA_PUBKEY_OUT"
    printf 'OK: wrote CA pubkey to %s\n' "$CA_PUBKEY_OUT"
  fi

  if [ -n "$K8S_CA_SECRET" ]; then
    local tmp
    tmp="$(mktemp)"
    printf '%s\n' "$pubkey" >"$tmp"
    kubectl_exec create secret generic "$K8S_CA_SECRET" \
      --namespace "$OPENBAO_NAMESPACE" \
      --from-file=ca_user.pub="$tmp" \
      --dry-run=client -o yaml | kubectl_exec apply -f -
    rm -f "$tmp"
    printf 'OK: K8s secret %s/%s updated\n' "$OPENBAO_NAMESPACE" "$K8S_CA_SECRET"
  fi
}

token="$(read_token)"
if [ -z "$token" ]; then
  echo "ERROR: empty token" >&2
  exit 1
fi

remote_bao "$token" status
enable_ssh_engine "$token"
ensure_default_issuer "$token"
apply_roles "$token"
write_policy "$token" warden-sign "$POLICY_DIR/warden-sign.hcl"
remote_bao "$token" list "${SSH_MOUNT}/roles"
export_ca_pubkey "$token"

cat <<NEXT

OpenBao SSH engine configuration applied.

Next steps:
  1. make openbao-verify-ssh
  2. railiance-infra: make bootstrap-ssh-ca SSH_CA_PUBKEY=<path>
  3. ops-warden: warden sign smoke (WP-0008 T2)
NEXT

if [ "$WARNINGS" -gt 0 ]; then
  printf '\nCompleted with %s warning(s).\n' "$WARNINGS"
fi