#!/usr/bin/env bash
# ============================================================
# Velocity-OS — Ingress Box: Air-Gap Transfer Agent
# Runs on a LAN-connected node (Raspberry Pi / VM).
# Polls ECR every 5 minutes for new signed images.
# Verifies cosign signature. Transfers to air-gapped workstation.
# Triggers K3s rolling restart on new image.
#
# Install as systemd service:
#   sudo cp poll_and_transfer.service /etc/systemd/system/
#   sudo systemctl enable --now poll_and_transfer
# ============================================================
set -euo pipefail

# ── Configuration ────────────────────────────────────────────
AWS_REGION="${AWS_REGION:-ap-south-1}"
AWS_ACCOUNT_ID="${AWS_ACCOUNT_ID:?Must set AWS_ACCOUNT_ID}"
ECR_REGISTRY="${AWS_ACCOUNT_ID}.dkr.ecr.${AWS_REGION}.amazonaws.com"
REGISTRY_PREFIX="velocity-os"
SERVICES=("core" "webos" "media-engine" "agents")

# Air-gapped workstation (LAN only — no internet)
WORKSTATION_IP="${WORKSTATION_IP:-192.168.1.100}"
WORKSTATION_USER="${WORKSTATION_USER:-ubuntu}"
WORKSTATION_SSH_KEY="${WORKSTATION_SSH_KEY:-/home/ingress/.ssh/velocity_workstation_ed25519}"

# State file: tracks last-transferred digest per service
STATE_DIR="/var/lib/velocity-ingress"
mkdir -p "${STATE_DIR}"

# Temp dir for image tarballs
TRANSFER_DIR="/tmp/velocity-transfer"
mkdir -p "${TRANSFER_DIR}"

# ── Functions ─────────────────────────────────────────────────
log() { echo "[$(date -u +%Y-%m-%dT%H:%M:%SZ)] $*"; }

get_latest_digest() {
  local repo="${REGISTRY_PREFIX}/$1"
  aws ecr describe-images \
    --repository-name "${repo}" \
    --image-ids imageTag=latest \
    --region "${AWS_REGION}" \
    --query 'imageDetails[0].imageDigest' \
    --output text 2>/dev/null || echo "NONE"
}

transfer_image() {
  local svc="$1"
  local digest="$2"
  local full_image="${ECR_REGISTRY}/${REGISTRY_PREFIX}/${svc}@${digest}"
  local tar_file="${TRANSFER_DIR}/${svc}.tar"

  log "  [${svc}] Pulling from ECR..."
  docker pull "${ECR_REGISTRY}/${REGISTRY_PREFIX}/${svc}:latest"

  log "  [${svc}] Verifying cosign signature..."
  cosign verify \
    --certificate-identity-regexp ".*" \
    --certificate-oidc-issuer-regexp ".*" \
    "${full_image}" || {
    log "  [${svc}] ERROR: Signature verification FAILED. Refusing transfer."
    return 1
  }

  log "  [${svc}] Saving image to tarball..."
  docker save "${ECR_REGISTRY}/${REGISTRY_PREFIX}/${svc}:latest" \
    -o "${tar_file}"

  log "  [${svc}] Transferring to workstation via SCP..."
  scp -i "${WORKSTATION_SSH_KEY}" \
    -o StrictHostKeyChecking=yes \
    "${tar_file}" \
    "${WORKSTATION_USER}@${WORKSTATION_IP}:/tmp/${svc}.tar"

  log "  [${svc}] Importing into K3s containerd + rolling restart..."
  ssh -i "${WORKSTATION_SSH_KEY}" \
    -o StrictHostKeyChecking=yes \
    "${WORKSTATION_USER}@${WORKSTATION_IP}" \
    "sudo k3s ctr images import /tmp/${svc}.tar && \
     sudo kubectl rollout restart deployment/${svc} -n velocity-os && \
     rm /tmp/${svc}.tar"

  # Record transferred digest
  echo "${digest}" > "${STATE_DIR}/${svc}.last_digest"
  log "  [${svc}] ✓ Transfer complete. Digest: ${digest}"
  rm -f "${tar_file}"
}

# ── Main poll loop ────────────────────────────────────────────
log "=== Velocity-OS Ingress Box polling ECR ==="

# Login to ECR (token expires every 12h; cron re-runs this)
aws ecr get-login-password --region "${AWS_REGION}" | \
  docker login --username AWS --password-stdin "${ECR_REGISTRY}"

for svc in "${SERVICES[@]}"; do
  log "[${svc}] Checking for updates..."

  CURRENT_DIGEST=$(get_latest_digest "${svc}")
  LAST_DIGEST=$(cat "${STATE_DIR}/${svc}.last_digest" 2>/dev/null || echo "NONE")

  if [[ "${CURRENT_DIGEST}" == "NONE" ]]; then
    log "  [${svc}] No image found in ECR. Skipping."
    continue
  fi

  if [[ "${CURRENT_DIGEST}" == "${LAST_DIGEST}" ]]; then
    log "  [${svc}] Up to date. No transfer needed."
    continue
  fi

  log "  [${svc}] New digest detected: ${CURRENT_DIGEST}"
  transfer_image "${svc}" "${CURRENT_DIGEST}" || \
    log "  [${svc}] Transfer FAILED. Will retry next cycle."
done

log "=== Poll cycle complete ==="