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

ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
REPO="${REPO:-Sybil-2}"
WORKLOAD="${WORKLOAD:-sybil-dev}"
CLUSTER="${CLUSTER:-viper}"
NAMESPACE="${NAMESPACE:-buzzert}"
QUOTA="${QUOTA:-research-api}"
SECRET="${SECRET:-sybil-dev-env}"
REMOTE_ROOT="${REMOTE_ROOT:-/root/code/$REPO}"
INGRESS_PORT="${INGRESS_PORT:-5173}"
CODEX_IMAGE="${CODEX_IMAGE:-openai.azurecr.io/rcall:crow-1b540b5e15f9-2026-01-22-2404-312}"
CODEX_SERVICE="${CODEX_SERVICE:-sybil-codex-devbox}"
CODEX_CONTAINER="${CODEX_CONTAINER:-codex-devbox}"
CODEX_SSH_SECRET="${CODEX_SSH_SECRET:-sybil-codex-ssh}"
CODEX_SSH_PORT="${CODEX_SSH_PORT:-2222}"
CODEX_WORKDIR="${CODEX_WORKDIR:-/workspace/sybil-codex}"
CODEX_SSH_KEY_PATH="${CODEX_SSH_KEY_PATH:-/run/secrets/sybil-codex-ssh/id_ed25519}"

echo "Syncing $REPO to $WORKLOAD on $CLUSTER/$NAMESPACE..."
brix create \
  --project-root "$(dirname "$ROOT")" \
  --path "$REPO=$ROOT" \
  --repositories "$REPO" \
  --include-untracked \
  --reconcile-policy continuous \
  -- \
  "cluster=$CLUSTER" \
  "namespace=$NAMESPACE" \
  "name=$WORKLOAD" \
  "pool.spec.quota=$QUOTA" \
  "pool.spec.size.resource=cpu" \
  "pool.spec.size.quantity=1" \
  "pool.spec.priority=low"

if [[ "$REPO" != "openai" ]]; then
  if kubectl --context "$CLUSTER" -n "$NAMESPACE" get git "$WORKLOAD" -o json 2>/dev/null | jq -e '.spec.repositories.openai?' >/dev/null; then
    kubectl --context "$CLUSTER" -n "$NAMESPACE" patch git "$WORKLOAD" --type=json -p='[{"op":"remove","path":"/spec/repositories/openai"}]'
  fi
fi

if [[ -n "${OPENAI_API_KEY:-}" && -n "${EXA_API_KEY:-}" ]]; then
  kubectl --context "$CLUSTER" -n "$NAMESPACE" create secret generic "$SECRET" \
    --from-literal=OPENAI_API_KEY="$OPENAI_API_KEY" \
    --from-literal=EXA_API_KEY="$EXA_API_KEY" \
    --dry-run=client -o yaml | kubectl --context "$CLUSTER" -n "$NAMESPACE" apply -f -
elif [[ -n "${OPENAI_API_KEY:-}" || -n "${EXA_API_KEY:-}" ]]; then
  echo "Only one API key env var is set; keeping existing Kubernetes secret $SECRET." >&2
fi

if [[ "$(kubectl --context "$CLUSTER" -n "$NAMESPACE" get secret "$SECRET" -o name 2>/dev/null || true)" == "" ]]; then
  echo "Secret $SECRET does not exist. Re-run with OPENAI_API_KEY and EXA_API_KEY set." >&2
  exit 1
fi

if [[ "$(kubectl --context "$CLUSTER" -n "$NAMESPACE" get secret "$CODEX_SSH_SECRET" -o name 2>/dev/null || true)" == "" ]]; then
  tmpdir="$(mktemp -d)"
  trap 'rm -rf "$tmpdir"' EXIT
  ssh-keygen -t ed25519 -N "" -C "sybil-codex@$WORKLOAD" -f "$tmpdir/id_ed25519" >/dev/null
  kubectl --context "$CLUSTER" -n "$NAMESPACE" create secret generic "$CODEX_SSH_SECRET" \
    --from-file=id_ed25519="$tmpdir/id_ed25519" \
    --from-file=id_ed25519.pub="$tmpdir/id_ed25519.pub" \
    --dry-run=client -o yaml | kubectl --context "$CLUSTER" -n "$NAMESPACE" apply -f -
fi

kubectl --context "$CLUSTER" -n "$NAMESPACE" patch secret "$SECRET" --type=merge -p "$(jq -n \
  --arg host "$CODEX_SERVICE" \
  --arg port "$CODEX_SSH_PORT" \
  --arg workdir "$CODEX_WORKDIR" \
  --arg keyPath "$CODEX_SSH_KEY_PATH" \
  '{
    stringData: {
      CHAT_CODEX_TOOL_ENABLED: "true",
      CHAT_CODEX_REMOTE_HOST: $host,
      CHAT_CODEX_REMOTE_USER: "root",
      CHAT_CODEX_REMOTE_PORT: $port,
      CHAT_CODEX_REMOTE_WORKDIR: $workdir,
      CHAT_CODEX_SSH_KEY_PATH: $keyPath,
      CHAT_CODEX_EXEC_TIMEOUT_MS: "600000",
      CHAT_SHELL_TOOL_ENABLED: "true"
    }
  }')"

ensure_codex_deployment() {
  kubectl --context "$CLUSTER" -n "$NAMESPACE" apply -f - <<KUBE_MANIFEST
apiVersion: apps/v1
kind: Deployment
metadata:
  name: $CODEX_SERVICE
  labels:
    app.kubernetes.io/name: $CODEX_SERVICE
    app.kubernetes.io/part-of: sybil-dev
spec:
  replicas: 1
  selector:
    matchLabels:
      app.kubernetes.io/name: $CODEX_SERVICE
  template:
    metadata:
      labels:
        app.kubernetes.io/name: $CODEX_SERVICE
        app.kubernetes.io/part-of: sybil-dev
    spec:
      tolerations:
        - key: openai.com/team
          operator: Exists
          effect: NoSchedule
      containers:
        - name: $CODEX_CONTAINER
          image: $CODEX_IMAGE
          imagePullPolicy: IfNotPresent
          command: ["/bin/bash", "-lc"]
          args:
            - |
              set -euo pipefail
              export DEBIAN_FRONTEND=noninteractive
              if ! command -v sshd >/dev/null 2>&1; then
                apt-get update
                apt-get install -y --no-install-recommends openssh-server ca-certificates git curl
                rm -rf /var/lib/apt/lists/*
              fi
              if ! command -v codex >/dev/null 2>&1; then
                npm install -g @openai/codex
              fi
              mkdir -p /run/sshd /root/.ssh "$CODEX_WORKDIR"
              cp /run/secrets/sybil-codex-ssh/id_ed25519.pub /root/.ssh/authorized_keys
              chmod 700 /root/.ssh
              chmod 600 /root/.ssh/authorized_keys
              {
                if [[ -n "\${OPENAI_API_KEY:-}" ]]; then
                  printf 'export OPENAI_API_KEY=%q\n' "\$OPENAI_API_KEY"
                fi
                if [[ -n "\${EXA_API_KEY:-}" ]]; then
                  printf 'export EXA_API_KEY=%q\n' "\$EXA_API_KEY"
                fi
              } >/root/.sybil-codex-env
              chmod 600 /root/.sybil-codex-env
              cat >/usr/local/bin/sybil-codex-ssh-command <<'SSH_COMMAND'
              #!/usr/bin/env bash
              set -euo pipefail
              if [[ -f /root/.sybil-codex-env ]]; then
                source /root/.sybil-codex-env
              fi
              if [[ -n "\${SSH_ORIGINAL_COMMAND:-}" ]]; then
                exec /bin/bash -lc "\$SSH_ORIGINAL_COMMAND"
              fi
              exec /bin/bash -l
              SSH_COMMAND
              chmod 700 /usr/local/bin/sybil-codex-ssh-command
              ssh-keygen -A
              cat >/tmp/sybil-codex-sshd_config <<SSHD_CONFIG
              Port $CODEX_SSH_PORT
              ListenAddress 0.0.0.0
              PermitRootLogin prohibit-password
              PasswordAuthentication no
              PubkeyAuthentication yes
              AcceptEnv OPENAI_API_KEY EXA_API_KEY
              AuthorizedKeysFile .ssh/authorized_keys
              PermitUserEnvironment no
              ForceCommand /usr/local/bin/sybil-codex-ssh-command
              AllowTcpForwarding no
              X11Forwarding no
              Subsystem sftp internal-sftp
              PidFile /run/sshd/sybil-codex.pid
              SSHD_CONFIG
              exec \$(command -v sshd) -D -e -f /tmp/sybil-codex-sshd_config
          ports:
            - name: ssh
              containerPort: $CODEX_SSH_PORT
          envFrom:
            - secretRef:
                name: $SECRET
          resources:
            requests:
              cpu: "1"
              memory: 4Gi
            limits:
              cpu: "1"
              memory: 4Gi
          volumeMounts:
            - name: sybil-codex-ssh
              mountPath: /run/secrets/sybil-codex-ssh
              readOnly: true
            - name: sybil-codex-workspace
              mountPath: /workspace
      volumes:
        - name: sybil-codex-ssh
          secret:
            secretName: $CODEX_SSH_SECRET
            defaultMode: 384
        - name: sybil-codex-workspace
          emptyDir: {}
---
apiVersion: v1
kind: Service
metadata:
  name: $CODEX_SERVICE
  labels:
    app.kubernetes.io/name: $CODEX_SERVICE
    app.kubernetes.io/part-of: sybil-dev
spec:
  type: ClusterIP
  selector:
    app.kubernetes.io/name: $CODEX_SERVICE
  ports:
    - name: ssh
      port: $CODEX_SSH_PORT
      targetPort: ssh
KUBE_MANIFEST

  kubectl --context "$CLUSTER" -n "$NAMESPACE" rollout status "deployment/$CODEX_SERVICE" --timeout=300s
}

ensure_codex_deployment

ensure_env_from() {
  local resource="$1"
  local name="$2"
  local jsonpath="$3"
  local patch_path="$4"
  local existing_env_from

  existing_env_from="$(kubectl --context "$CLUSTER" -n "$NAMESPACE" get "$resource" "$name" \
    -o "jsonpath=$jsonpath" 2>/dev/null || true)"

  if [[ "$existing_env_from" == *"$SECRET"* ]]; then
    return
  fi

  local patch
  if [[ -z "$existing_env_from" ]]; then
    patch="[{\"op\":\"add\",\"path\":\"$patch_path\",\"value\":[{\"secretRef\":{\"name\":\"$SECRET\"}}]}]"
  else
    patch="[{\"op\":\"add\",\"path\":\"$patch_path/-\",\"value\":{\"secretRef\":{\"name\":\"$SECRET\"}}}]"
  fi
  kubectl --context "$CLUSTER" -n "$NAMESPACE" patch "$resource" "$name" --type=json -p "$patch"
}

ensure_env_from \
  workload \
  "$WORKLOAD" \
  "{.spec.pools[0].template.template.spec.workers.template.spec.containers[0].envFrom}" \
  "/spec/pools/0/template/template/spec/workers/template/spec/containers/0/envFrom"

ensure_env_from \
  pool \
  "$WORKLOAD" \
  "{.spec.workers.template.spec.containers[0].envFrom}" \
  "/spec/workers/template/spec/containers/0/envFrom"

codex_ssh_volume_json="$(jq -n --arg secret "$CODEX_SSH_SECRET" '{
  name: "sybil-codex-ssh",
  secret: {
    secretName: $secret,
    defaultMode: 384
  }
}')"
codex_main_mount_json="$(jq -n '{name: "sybil-codex-ssh", mountPath: "/run/secrets/sybil-codex-ssh", readOnly: true}')"

patch_pool_codex_mount() {
  local pool_json containers volumes patch
  pool_json="$(kubectl --context "$CLUSTER" -n "$NAMESPACE" get pool "$WORKLOAD" -o json)"
  containers="$(jq \
    --arg sidecarName "$CODEX_CONTAINER" \
    --argjson mainMount "$codex_main_mount_json" \
    '
    def ensure_mount($mount):
      .volumeMounts = (((.volumeMounts // []) | map(select(.name != $mount.name))) + [$mount]);
    (.spec.workers.template.spec.containers // [])
    | map(if .name == "main" then ensure_mount($mainMount) else . end)
    | map(select(.name != $sidecarName))
    ' <<<"$pool_json")"
  volumes="$(jq \
    --argjson sshVolume "$codex_ssh_volume_json" \
    '
    (.spec.workers.template.spec.volumes // [])
    | map(select(.name != $sshVolume.name and .name != "sybil-codex-workspace"))
    + [$sshVolume]
    ' <<<"$pool_json")"
  patch="$(jq -n --argjson containers "$containers" --argjson volumes "$volumes" '{
    spec: {
      workers: {
        template: {
          spec: {
            containers: $containers,
            volumes: $volumes
          }
        }
      }
    }
  }')"
  kubectl --context "$CLUSTER" -n "$NAMESPACE" patch pool "$WORKLOAD" --type=merge -p "$patch"
}

patch_workload_codex_mount() {
  local workload_json pools patch
  workload_json="$(kubectl --context "$CLUSTER" -n "$NAMESPACE" get workload "$WORKLOAD" -o json)"
  pools="$(jq \
    --arg groupName "$WORKLOAD" \
    --arg sidecarName "$CODEX_CONTAINER" \
    --argjson mainMount "$codex_main_mount_json" \
    --argjson sshVolume "$codex_ssh_volume_json" \
    '
    def ensure_mount($mount):
      .volumeMounts = (((.volumeMounts // []) | map(select(.name != $mount.name))) + [$mount]);
    .spec.pools
    | map(
        if .groupName == $groupName then
          .template.template.spec.workers.template.spec.containers =
            ((.template.template.spec.workers.template.spec.containers // [])
              | map(if .name == "main" then ensure_mount($mainMount) else . end)
              | map(select(.name != $sidecarName)))
          | .template.template.spec.workers.template.spec.volumes =
            ((.template.template.spec.workers.template.spec.volumes // [])
              | map(select(.name != $sshVolume.name and .name != "sybil-codex-workspace"))
              + [$sshVolume])
        else
          .
        end
      )
    ' <<<"$workload_json")"
  patch="$(jq -n --argjson pools "$pools" '{spec: {pools: $pools}}')"
  kubectl --context "$CLUSTER" -n "$NAMESPACE" patch workload "$WORKLOAD" --type=merge -p "$patch"
}

patch_workload_codex_mount
patch_pool_codex_mount

echo "Waiting for $WORKLOAD-0 to run the current pool revision..."
deadline=$((SECONDS + 300))
while true; do
  pool_revision="$(kubectl --context "$CLUSTER" -n "$NAMESPACE" get pool "$WORKLOAD" -o jsonpath='{.status.revision}' 2>/dev/null || true)"
  pod_revision="$(kubectl --context "$CLUSTER" -n "$NAMESPACE" get pod "$WORKLOAD-0" -o jsonpath='{.metadata.labels.brix\.openai\.com/revision}' 2>/dev/null || true)"
  pod_ready="$(kubectl --context "$CLUSTER" -n "$NAMESPACE" get pod "$WORKLOAD-0" -o jsonpath='{.status.conditions[?(@.type=="Ready")].status}' 2>/dev/null || true)"
  pod_deletion_timestamp="$(kubectl --context "$CLUSTER" -n "$NAMESPACE" get pod "$WORKLOAD-0" -o jsonpath='{.metadata.deletionTimestamp}' 2>/dev/null || true)"

  if [[ -n "$pool_revision" && "$pod_revision" == "$pool_revision" && "$pod_ready" == "True" && -z "$pod_deletion_timestamp" ]]; then
    break
  fi

  if (( SECONDS >= deadline )); then
    echo "Timed out waiting for current ready pod. pool_revision=$pool_revision pod_revision=$pod_revision pod_ready=$pod_ready deleting=$pod_deletion_timestamp" >&2
    exit 1
  fi

  sleep 5
done

brix run --clusters "$CLUSTER" --namespace "$NAMESPACE" --pools "$WORKLOAD" --dir "$REMOTE_ROOT" -- ./scripts/devbox-start.sh

ssh_check="ssh -n -i '$CODEX_SSH_KEY_PATH' -o BatchMode=yes -o StrictHostKeyChecking=accept-new -o UserKnownHostsFile=/tmp/sybil-codex-known-hosts -p '$CODEX_SSH_PORT' root@'$CODEX_SERVICE' true"
brix run --clusters "$CLUSTER" --namespace "$NAMESPACE" --pools "$WORKLOAD" --dir "$REMOTE_ROOT" -- /bin/bash -lc "$ssh_check"

url="$(brix ingress create --cluster "$CLUSTER" --namespace "$NAMESPACE" --pod "$WORKLOAD-0" --ports "web:$INGRESS_PORT" --root web --timeout 300 "$WORKLOAD" | tail -1)"

echo "Devbox update complete: $url"