#!/usr/bin/env bash
# neuron-self-improve.sh — self-improvement loop for Neuron
#
# Manages code experiments across alpha/beta/gamma zones and promotes winners to prod.
# All cluster changes go through git → Argo CD. No direct kubectl mutations.
#
# Usage:
#   ./neuron-self-improve.sh snapshot <slot>          # copy prod DB to slot
#   ./neuron-self-improve.sh deploy <slot> <sha>      # point slot at a specific image
#   ./neuron-self-improve.sh promote <slot>           # copy slot's image to prod
#   ./neuron-self-improve.sh status                   # show what's running in each zone
#   ./neuron-self-improve.sh rollback <slot> <sha>    # revert slot to a previous image

set -euo pipefail

INFRA_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
REGISTRY="registry.neuralplatform.ai/neural-platform/neuron"
KUBECONFIG="${KUBECONFIG:-$HOME/.kube/legion-config}"

# ── helpers ───────────────────────────────────────────────────────────────────

slot_manifest() {
  local slot=$1
  if [[ "$slot" == "prod" ]]; then
    echo "$INFRA_DIR/servers/legion/apps/neuron.yaml"
  else
    echo "$INFRA_DIR/servers/legion/apps/neuron-${slot}.yaml"
  fi
}

current_sha() {
  local slot=$1
  grep "image:" "$(slot_manifest "$slot")" | grep -oP '[a-f0-9]{40}' | head -1
}

set_image() {
  local slot=$1 sha=$2
  local manifest
  manifest="$(slot_manifest "$slot")"
  sed -i "s|${REGISTRY}:.*|${REGISTRY}:${sha}|" "$manifest"
  echo "  set $slot → $sha"
}

git_push() {
  local msg=$1
  cd "$INFRA_DIR"
  git add -A
  git commit -m "$msg"
  git push
  echo "  pushed: $msg"
}

wait_healthy() {
  local slot=$1
  local app="neuron"
  [[ "$slot" != "prod" ]] && app="neuron-${slot}"
  echo "  waiting for $slot to be healthy..."
  local i=0
  while [[ $i -lt 30 ]]; do
    local ready
    ready=$(KUBECONFIG="$KUBECONFIG" kubectl get pods -n neuron -l "app=${app}" \
      --no-headers 2>/dev/null | grep "1/1" | wc -l)
    [[ "$ready" -ge 1 ]] && echo "  $slot is healthy" && return 0
    sleep 10
    ((i++))
  done
  echo "  WARNING: $slot did not become healthy within 5 minutes"
  return 1
}

# ── commands ──────────────────────────────────────────────────────────────────

cmd_snapshot() {
  local slot=${1:?Usage: snapshot <slot>}
  echo "==> Snapshotting prod DB into $slot..."
  KUBECONFIG="$KUBECONFIG" kubectl apply -f - << EOF
apiVersion: batch/v1
kind: Job
metadata:
  name: snapshot-prod-to-${slot}
  namespace: neuron
spec:
  ttlSecondsAfterFinished: 300
  template:
    spec:
      restartPolicy: Never
      containers:
        - name: snapshot
          image: keinos/sqlite3:latest
          command: ["/bin/sh", "-c"]
          args:
            - sqlite3 /src/neuron.db ".backup /dst/neuron.db" && echo "Done" && ls -lh /dst/neuron.db
          volumeMounts:
            - name: src
              mountPath: /src
              readOnly: true
            - name: dst
              mountPath: /dst
      volumes:
        - name: src
          persistentVolumeClaim:
            claimName: neuron-data
        - name: dst
          persistentVolumeClaim:
            claimName: neuron-${slot}-data
EOF
  echo "  waiting for snapshot job..."
  KUBECONFIG="$KUBECONFIG" kubectl wait job "snapshot-prod-to-${slot}" \
    -n neuron --for=condition=complete --timeout=120s
  echo "  snapshot complete"
}

cmd_deploy() {
  local slot=${1:?Usage: deploy <slot> <sha>}
  local sha=${2:?Usage: deploy <slot> <sha>}
  echo "==> Deploying $sha to $slot..."
  set_image "$slot" "$sha"
  git_push "chore(neuron): deploy ${sha:0:12} to $slot"
  wait_healthy "$slot"
  echo "==> $slot is live at https://neuron-${slot}.neuralplatform.ai"
}

cmd_promote() {
  local slot=${1:?Usage: promote <slot>}
  local sha
  sha=$(current_sha "$slot")
  [[ -z "$sha" ]] && echo "ERROR: could not read SHA from $slot" && exit 1
  echo "==> Promoting $slot ($sha) to prod..."
  set_image "prod" "$sha"
  # Also stamp the :stable registry tag (pull and re-tag)
  git_push "chore(neuron): promote $slot (${sha:0:12}) to prod"
  wait_healthy "prod"
  echo "==> prod is now running $sha"
  echo "==> prod endpoint: https://neuron.neuralplatform.ai"
}

cmd_rollback() {
  local slot=${1:?Usage: rollback <slot> <sha>}
  local sha=${2:?Usage: rollback <slot> <sha>}
  echo "==> Rolling $slot back to $sha..."
  set_image "$slot" "$sha"
  git_push "chore(neuron): rollback $slot to ${sha:0:12}"
  wait_healthy "$slot"
  echo "==> $slot rolled back"
}

cmd_status() {
  echo "==> Neuron slot status"
  echo ""
  for slot in prod alpha beta gamma; do
    local app="neuron"
    [[ "$slot" != "prod" ]] && app="neuron-${slot}"
    local sha
    sha=$(current_sha "$slot" 2>/dev/null || echo "unknown")
    local ready
    ready=$(KUBECONFIG="$KUBECONFIG" kubectl get pods -n neuron -l "app=${app}" \
      --no-headers 2>/dev/null | grep "1/1" | wc -l | tr -d ' ')
    printf "  %-8s  %s  pods_ready=%s\n" "$slot" "${sha:0:12}" "$ready"
  done
  echo ""
  echo "  Endpoints:"
  echo "    prod:  https://neuron.neuralplatform.ai"
  echo "    alpha: https://neuron-alpha.neuralplatform.ai"
  echo "    beta:  https://neuron-beta.neuralplatform.ai"
  echo "    gamma: https://neuron-gamma.neuralplatform.ai"
}

# ── self-improvement loop (autonomous) ────────────────────────────────────────
# When called with no args, prints the process for an autonomous agent to follow.
cmd_loop_process() {
  cat << 'PROCESS'
NEURON SELF-IMPROVEMENT LOOP
=============================

Each iteration:

1. IDENTIFY — find an improvement opportunity in the codebase.
   Sources: Neuron backlog, observed failures, performance gaps, code review.

2. IMPLEMENT — make the change on a git branch in neural-platform/neuron.
   Commit message should describe the hypothesis being tested.

3. CI — push the branch; CI runs lint/test/build and auto-deploys to alpha.
   Wait for CI to complete and alpha pod to be healthy.

4. SNAPSHOT — before the next experiment touches a slot, snapshot prod DB:
     ./scripts/neuron-self-improve.sh snapshot <slot>
   This ensures the experiment runs against real data without touching prod.

5. EXPERIMENT — run the experiment against the alpha endpoint:
     https://neuron-alpha.neuralplatform.ai
   Use real workloads. Compare against prod or other slots.
   Multiple experiments can run simultaneously across alpha/beta/gamma.

6. EVALUATE — compare results. Criteria:
   - Does it behave correctly?
   - Does it perform better (latency, quality, fewer errors)?
   - Does it handle edge cases prod handles?
   Schema changes are acceptable — they become migrations against prod on promote.

7. PROMOTE WINNER — if alpha wins:
     ./scripts/neuron-self-improve.sh promote alpha
   This updates prod's image SHA in git → Argo CD deploys → prod is updated.
   The old prod SHA is still in the registry — rollback is always available.

8. ROLLBACK — if something goes wrong:
     ./scripts/neuron-self-improve.sh rollback prod <previous-sha>

9. LOOP — go to step 1. Alpha is now free for the next experiment.

RULES
------
- Never skip the snapshot step before running an experiment against a slot.
- Prod data is never lost — slots only get copies.
- Schema changes in experiment code → Alembic auto-migrates the snapshot DB.
  On promote, the migration runs against prod. Always verify migrations are
  additive/non-destructive before promoting.
- All three experiment zones can run simultaneously with different hypotheses.
- The registry retains every SHA — any version can be deployed to any slot.
- This script handles git push; Argo CD handles k8s — never kubectl apply.

PROCESS
}

# ── dispatch ──────────────────────────────────────────────────────────────────

case "${1:-loop}" in
  snapshot) cmd_snapshot "${@:2}" ;;
  deploy)   cmd_deploy   "${@:2}" ;;
  promote)  cmd_promote  "${@:2}" ;;
  rollback) cmd_rollback "${@:2}" ;;
  status)   cmd_status ;;
  loop)     cmd_loop_process ;;
  *)        echo "Usage: $0 snapshot|deploy|promote|rollback|status|loop" && exit 1 ;;
esac