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

MODE="${1:-check}"
shift || true

NAMESPACE="${AWOOOI_NAMESPACE:-awoooi-prod}"
BROKER_DEPLOYMENT="${AWOOOI_DB_BOOTSTRAP_BROKER_DEPLOYMENT:-awoooi-ansible-executor-broker}"
BROKER_CONTAINER="${AWOOOI_DB_BOOTSTRAP_BROKER_CONTAINER:-broker}"
POSTGRES_SSH_TARGET="${AWOOOI_DB_BOOTSTRAP_POSTGRES_SSH_TARGET:-ollama@192.168.0.188}"
POSTGRES_CONTAINER="${AWOOOI_DB_BOOTSTRAP_POSTGRES_CONTAINER:-k3s-postgres-recovery}"
POSTGRES_DATABASE="${AWOOOI_DB_BOOTSTRAP_DATABASE:-awoooi_prod}"
POSTGRES_HOST="${AWOOOI_DB_BOOTSTRAP_DATABASE_HOST:-192.168.0.188}"
POSTGRES_PORT="${AWOOOI_DB_BOOTSTRAP_DATABASE_PORT:-5432}"
SHARED_APP_ROLE="${AWOOOI_DB_BOOTSTRAP_SHARED_ROLE:-awoooi}"
RLS_APP_ROLE="${AWOOOI_DB_BOOTSTRAP_RLS_ROLE:-awooop_app}"
SSH_KEY_PATH="${AWOOOI_DB_BOOTSTRAP_SSH_KEY_PATH:-/run/secrets/ssh_mcp_key}"
KNOWN_HOSTS_PATH="${AWOOOI_DB_BOOTSTRAP_KNOWN_HOSTS_PATH:-/etc/ssh-mcp/known_hosts}"

case "$MODE" in
  check|apply|verify|disable) ;;
  *)
    echo "usage: $0 check|apply|verify|disable [api|worker|broker ...]" >&2
    exit 2
    ;;
esac

KUBECTL=(kubectl)
if [ "${AWOOOI_KUBECTL_SUDO:-false}" = "true" ]; then
  KUBECTL=(sudo kubectl)
fi
if [ -n "${AWOOOI_KUBECONFIG:-}" ]; then
  KUBECTL+=(--kubeconfig="${AWOOOI_KUBECONFIG}")
fi
if [ -n "${AWOOOI_KUBERNETES_SERVER:-}" ]; then
  KUBECTL+=(--server="${AWOOOI_KUBERNETES_SERVER}")
fi

kube() {
  "${KUBECTL[@]}" "$@"
}

workload_role() {
  case "$1" in
    api) printf 'awoooi_api_runtime' ;;
    worker) printf 'awoooi_worker_runtime' ;;
    broker) printf 'awoooi_broker_runtime' ;;
    *) return 1 ;;
  esac
}

workload_secret() {
  case "$1" in
    api) printf 'awoooi-api-db' ;;
    worker) printf 'awoooi-worker-db' ;;
    broker) printf 'awoooi-broker-db' ;;
    *) return 1 ;;
  esac
}

workload_connection_limit() {
  case "$1" in
    api) printf '12' ;;
    worker) printf '8' ;;
    broker) printf '4' ;;
    *) return 1 ;;
  esac
}

workload_required_connection_budget() {
  case "$1" in
    api) printf '8' ;;
    worker) printf '5' ;;
    broker) printf '2' ;;
    *) return 1 ;;
  esac
}

workload_deployment() {
  case "$1" in
    api) printf 'awoooi-api|api' ;;
    worker) printf 'awoooi-worker|worker' ;;
    broker) printf 'awoooi-ansible-executor-broker|broker' ;;
    *) return 1 ;;
  esac
}

selected_workloads=("$@")
if [ "${#selected_workloads[@]}" -eq 0 ]; then
  selected_workloads=(api worker broker)
fi
for workload in "${selected_workloads[@]}"; do
  workload_role "$workload" >/dev/null || {
    echo "invalid workload id: $workload" >&2
    exit 2
  }
done

admin_psql() {
  kube exec -i "deployment/${BROKER_DEPLOYMENT}" -n "$NAMESPACE" \
    -c "$BROKER_CONTAINER" -- \
    ssh -i "$SSH_KEY_PATH" \
      -o BatchMode=yes \
      -o ConnectTimeout=10 \
      -o StrictHostKeyChecking=yes \
      -o UserKnownHostsFile="$KNOWN_HOSTS_PATH" \
      "$POSTGRES_SSH_TARGET" \
      "docker exec -i -u postgres ${POSTGRES_CONTAINER} psql -h /tmp -U postgres -d ${POSTGRES_DATABASE} -v ON_ERROR_STOP=1 -Atq"
}

require_control_path() {
  kube get "deployment/${BROKER_DEPLOYMENT}" -n "$NAMESPACE" >/dev/null
  kube exec "deployment/${BROKER_DEPLOYMENT}" -n "$NAMESPACE" \
    -c "$BROKER_CONTAINER" -- test -s "$SSH_KEY_PATH"
  kube exec "deployment/${BROKER_DEPLOYMENT}" -n "$NAMESPACE" \
    -c "$BROKER_CONTAINER" -- test -s "$KNOWN_HOSTS_PATH"
  test "$(printf 'SELECT 1;\n' | admin_psql | tail -n 1)" = "1"
  role_count="$(printf \
    "SELECT count(*) FROM pg_roles WHERE rolname IN ('%s','%s');\n" \
    "$SHARED_APP_ROLE" "$RLS_APP_ROLE" | admin_psql | tail -n 1)"
  test "$role_count" = "2" || {
    echo "required shared database roles are unavailable" >&2
    exit 1
  }
}

role_state() {
  role="$1"
  printf "SELECT concat(rolcanlogin, '|', rolconnlimit, '|', pg_has_role('%s', '%s', 'member'), '|', pg_has_role('%s', '%s', 'member')) FROM pg_roles WHERE rolname = '%s';\n" \
    "$role" "$SHARED_APP_ROLE" "$role" "$RLS_APP_ROLE" "$role" \
    | admin_psql | tail -n 1
}

disable_role() {
  role="$1"
  cat <<SQL | admin_psql >/dev/null
DO \$disable\$
BEGIN
  IF EXISTS (SELECT 1 FROM pg_roles WHERE rolname = '${role}') THEN
    EXECUTE format('ALTER ROLE %I NOLOGIN', '${role}');
  END IF;
END
\$disable\$;
SQL
}

ensure_existing_role_contract() {
  role="$1"
  connection_limit="$2"
  cat <<SQL | admin_psql >/dev/null
DO \$contract\$
BEGIN
  IF NOT EXISTS (SELECT 1 FROM pg_roles WHERE rolname = '${role}') THEN
    RAISE EXCEPTION 'workload database role missing';
  END IF;
  EXECUTE format(
    'ALTER ROLE %I LOGIN INHERIT NOSUPERUSER NOCREATEDB NOCREATEROLE NOREPLICATION CONNECTION LIMIT %s',
    '${role}', ${connection_limit}
  );
END
\$contract\$;
GRANT ${SHARED_APP_ROLE} TO ${role};
GRANT ${RLS_APP_ROLE} TO ${role};
SQL
}

create_role_and_secret() {
  workload="$1"
  role="$2"
  secret="$3"
  connection_limit="$4"
  password="$(openssl rand -hex 32)"
  test "${#password}" -eq 64

  cat <<SQL | admin_psql >/dev/null
DO \$bootstrap\$
BEGIN
  IF EXISTS (SELECT 1 FROM pg_roles WHERE rolname = '${role}') THEN
    EXECUTE format(
      'ALTER ROLE %I LOGIN INHERIT NOSUPERUSER NOCREATEDB NOCREATEROLE NOREPLICATION CONNECTION LIMIT %s PASSWORD %L',
      '${role}', ${connection_limit}, '${password}'
    );
  ELSE
    EXECUTE format(
      'CREATE ROLE %I LOGIN INHERIT NOSUPERUSER NOCREATEDB NOCREATEROLE NOREPLICATION CONNECTION LIMIT %s PASSWORD %L',
      '${role}', ${connection_limit}, '${password}'
    );
  END IF;
END
\$bootstrap\$;
GRANT ${SHARED_APP_ROLE} TO ${role};
GRANT ${RLS_APP_ROLE} TO ${role};
SQL

  database_url="postgresql+asyncpg://${role}:${password}@${POSTGRES_HOST}:${POSTGRES_PORT}/${POSTGRES_DATABASE}?ssl=disable"
  if ! printf '%s' "$database_url" \
    | kube create secret generic "$secret" -n "$NAMESPACE" \
        --from-file=DATABASE_URL=/dev/stdin \
        --dry-run=client -o yaml \
    | kube apply -f - >/dev/null; then
    disable_role "$role"
    unset database_url password
    echo "workload database Secret apply failed; role was disabled" >&2
    return 1
  fi
  if ! kube label secret "$secret" -n "$NAMESPACE" \
    app.kubernetes.io/managed-by=awoooi-workload-db-identity-bootstrap \
    awoooi.wooo.work/database-workload="$workload" --overwrite >/dev/null; then
    kube delete secret "$secret" -n "$NAMESPACE" \
      --ignore-not-found=true >/dev/null 2>&1 || true
    disable_role "$role"
    unset database_url password
    echo "workload database Secret label failed; new Secret removed and role disabled" >&2
    return 1
  fi
  unset database_url password
}

preflight_workload() {
  workload="$1"
  secret="$2"
  connection_limit="$3"
  required_connection_budget="$4"
  deployment_spec="$(workload_deployment "$workload")"
  deployment="${deployment_spec%%|*}"
  container="${deployment_spec#*|}"
  image="$(kube get "deployment/${deployment}" -n "$NAMESPACE" \
    -o "jsonpath={.spec.template.spec.containers[?(@.name=='${container}')].image}")"
  test -n "$image"
  suffix="${AWOOOI_RUN_ID:-manual}-${RANDOM}"
  pod_name="awoooi-db-preflight-${workload}-${suffix}"
  pod_name="${pod_name:0:63}"

  cleanup_preflight() {
    kube delete pod "$pod_name" -n "$NAMESPACE" \
      --ignore-not-found=true --wait=false >/dev/null 2>&1 || true
  }
  cleanup_preflight
  cat <<YAML | kube apply -f - >/dev/null
apiVersion: v1
kind: Pod
metadata:
  name: ${pod_name}
  namespace: ${NAMESPACE}
  labels:
    system: awoooi
    component: workload-db-preflight
    awoooi.wooo.work/database-workload: ${workload}
spec:
  restartPolicy: Never
  automountServiceAccountToken: false
  containers:
    - name: preflight
      image: ${image}
      imagePullPolicy: IfNotPresent
      env:
        - name: DATABASE_URL
          valueFrom:
            secretKeyRef:
              name: ${secret}
              key: DATABASE_URL
        - name: EXPECTED_CONNECTION_LIMIT
          value: "${connection_limit}"
        - name: REQUIRED_CONCURRENT_CONNECTIONS
          value: "${required_connection_budget}"
        - name: WORKLOAD_ID
          value: "${workload}"
      command: ["python", "-c"]
      args:
        - |
          import asyncio
          import hashlib
          import json
          import os

          from sqlalchemy import text
          from sqlalchemy.ext.asyncio import create_async_engine
          from sqlalchemy.pool import NullPool


          async def main() -> None:
              engine = create_async_engine(
                  os.environ["DATABASE_URL"], poolclass=NullPool
              )
              try:
                  async with engine.connect() as conn:
                      identity = (
                          await conn.execute(text("""
                              SELECT current_user AS role_name, rolconnlimit
                              FROM pg_roles WHERE rolname = current_user
                          """))
                      ).mappings().one()
                      for table_name in (
                          "incidents",
                          "knowledge_entries",
                          "awooop_run_state",
                      ):
                          await conn.execute(
                              text(f'SELECT count(*) FROM "{table_name}"')
                          )
                      table_gap_count = int((await conn.execute(text("""
                          SELECT count(*)
                          FROM pg_class relation
                          JOIN pg_namespace namespace
                            ON namespace.oid = relation.relnamespace
                          WHERE namespace.nspname = 'public'
                            AND relation.relkind IN ('r', 'p', 'v', 'm')
                            AND (
                              (has_table_privilege('awoooi', relation.oid, 'SELECT')
                               AND NOT has_table_privilege(current_user, relation.oid, 'SELECT'))
                              OR (has_table_privilege('awoooi', relation.oid, 'INSERT')
                                  AND NOT has_table_privilege(current_user, relation.oid, 'INSERT'))
                              OR (has_table_privilege('awoooi', relation.oid, 'UPDATE')
                                  AND NOT has_table_privilege(current_user, relation.oid, 'UPDATE'))
                              OR (has_table_privilege('awoooi', relation.oid, 'DELETE')
                                  AND NOT has_table_privilege(current_user, relation.oid, 'DELETE'))
                            )
                      """))).scalar_one())
                      sequence_gap_count = int((await conn.execute(text("""
                          SELECT count(*)
                          FROM pg_class sequence
                          JOIN pg_namespace namespace
                            ON namespace.oid = sequence.relnamespace
                          WHERE namespace.nspname = 'public'
                            AND sequence.relkind = 'S'
                            AND (
                              (has_sequence_privilege('awoooi', sequence.oid, 'USAGE')
                               AND NOT has_sequence_privilege(current_user, sequence.oid, 'USAGE'))
                              OR (has_sequence_privilege('awoooi', sequence.oid, 'SELECT')
                                  AND NOT has_sequence_privilege(current_user, sequence.oid, 'SELECT'))
                              OR (has_sequence_privilege('awoooi', sequence.oid, 'UPDATE')
                                  AND NOT has_sequence_privilege(current_user, sequence.oid, 'UPDATE'))
                            )
                      """))).scalar_one())

                  required_concurrency = int(
                      os.environ["REQUIRED_CONCURRENT_CONNECTIONS"]
                  )

                  async def probe_connection() -> None:
                      async with engine.connect() as probe:
                          await probe.execute(text("SELECT pg_sleep(0.25)"))

                  await asyncio.gather(*(
                      probe_connection()
                      for _ in range(required_concurrency)
                  ))
              finally:
                  await engine.dispose()

              expected_limit = int(os.environ["EXPECTED_CONNECTION_LIMIT"])
              observed_limit = int(identity["rolconnlimit"])
              if observed_limit != expected_limit:
                  raise RuntimeError("connection_limit_mismatch")
              if table_gap_count or sequence_gap_count:
                  raise RuntimeError("database_privilege_gap")
              print(json.dumps({
                  "status": "verified",
                  "workload": os.environ["WORKLOAD_ID"],
                  "role_fingerprint": hashlib.sha256(
                      str(identity["role_name"]).encode("utf-8")
                  ).hexdigest(),
                  "connection_limit": observed_limit,
                  "concurrent_connection_probe_count": required_concurrency,
                  "representative_read_count": 3,
                  "table_privilege_gap_count": table_gap_count,
                  "sequence_privilege_gap_count": sequence_gap_count,
                  "secret_value_exposed": False,
                  "secret_value_logged": False,
              }, separators=(",", ":")))


          try:
              asyncio.run(main())
          except Exception as exc:
              print(json.dumps({
                  "status": "failed",
                  "workload": os.environ.get("WORKLOAD_ID", "unknown"),
                  "error_type": type(exc).__name__,
                  "secret_value_exposed": False,
                  "secret_value_logged": False,
              }, separators=(",", ":")))
              raise SystemExit(1)
YAML

  phase=""
  for _ in $(seq 1 60); do
    phase="$(kube get pod "$pod_name" -n "$NAMESPACE" \
      -o jsonpath='{.status.phase}' 2>/dev/null || true)"
    case "$phase" in
      Succeeded) break ;;
      Failed)
        kube logs "$pod_name" -n "$NAMESPACE" -c preflight || true
        cleanup_preflight
        return 1
        ;;
    esac
    sleep 2
  done
  if [ "$phase" != "Succeeded" ]; then
    echo "workload database preflight timed out: $workload" >&2
    cleanup_preflight
    return 1
  fi
  kube logs "$pod_name" -n "$NAMESPACE" -c preflight
  cleanup_preflight
}

require_control_path

if [ "$MODE" = "disable" ]; then
  for workload in "${selected_workloads[@]}"; do
    disable_role "$(workload_role "$workload")"
    echo "workload=$workload role_login_enabled=false"
  done
  exit 0
fi

new_workloads=()
rollback_new_workloads() {
  for workload in "${new_workloads[@]}"; do
    secret="$(workload_secret "$workload")"
    role="$(workload_role "$workload")"
    kube delete secret "$secret" -n "$NAMESPACE" \
      --ignore-not-found=true >/dev/null 2>&1 || true
    disable_role "$role" >/dev/null 2>&1 || true
    echo "workload=$workload bootstrap_rollback=secret_removed_role_disabled"
  done
}
if [ "$MODE" = "apply" ]; then
  trap 'status=$?; if [ "$status" -ne 0 ]; then rollback_new_workloads; fi; exit "$status"' EXIT
fi

ready=true
for workload in "${selected_workloads[@]}"; do
  role="$(workload_role "$workload")"
  secret="$(workload_secret "$workload")"
  connection_limit="$(workload_connection_limit "$workload")"
  required_connection_budget="$(workload_required_connection_budget "$workload")"
  secret_present=false
  if kube get secret "$secret" -n "$NAMESPACE" >/dev/null 2>&1; then
    secret_present=true
  fi
  state="$(role_state "$role")"
  role_present=false
  if [ -n "$state" ]; then
    role_present=true
  fi

  if [ "$MODE" = "apply" ]; then
    if [ "$secret_present" = "true" ]; then
      test "$role_present" = "true" || {
        echo "workload=$workload secret_present=true role_present=false" >&2
        exit 1
      }
      ensure_existing_role_contract "$role" "$connection_limit"
    else
      create_role_and_secret "$workload" "$role" "$secret" "$connection_limit"
      new_workloads+=("$workload")
      secret_present=true
      role_present=true
    fi
    state="$(role_state "$role")"
  fi

  expected_state="t|${connection_limit}|t|t"
  state_ready=false
  if [ "$secret_present" = "true" ] && [ "$state" = "$expected_state" ]; then
    state_ready=true
  else
    ready=false
  fi
  echo "workload=$workload secret_ref=$secret secret_present=$secret_present role_present=$role_present contract_ready=$state_ready connection_limit=$connection_limit required_connection_budget=$required_connection_budget"

  if [ "$MODE" = "apply" ] || [ "$MODE" = "verify" ]; then
    preflight_workload "$workload" "$secret" "$connection_limit" \
      "$required_connection_budget"
  fi
done

if [ "$ready" != "true" ]; then
  echo "WORKLOAD_DB_IDENTITY_BOOTSTRAP_READY=0"
  exit 1
fi
trap - EXIT
echo "WORKLOAD_DB_IDENTITY_BOOTSTRAP_READY=1"