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

ROOT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
ENV_FILE="${ROOT_DIR}/.env"

if [ -f "${ENV_FILE}" ]; then
    set -a
    # shellcheck disable=SC1090
    . "${ENV_FILE}"
    set +a
fi

IMAGE="${DOCKER_USERNAME:-bharathkumar192}/variant-worker:v2"
DISK="${DISK:-20}"
WORKER_PREFIX="${WORKER_PREFIX:-variant}"
DRY_RUN="false"

# Variant-specific tunables
VARIANT_MAX_JOBS="${VARIANT_MAX_JOBS:-0}"
VARIANT_BATCH_SIZE="${VARIANT_BATCH_SIZE:-10}"
VARIANT_CONCURRENT_REQUESTS="${VARIANT_CONCURRENT_REQUESTS:-20}"
VARIANT_PACK_TARGET_VIDEOS="${VARIANT_PACK_TARGET_VIDEOS:-50}"
VARIANT_PACK_TARGET_ROWS="${VARIANT_PACK_TARGET_ROWS:-5000}"

usage() {
    cat <<'EOF'
Usage:
  ./scripts/deploy_variant.sh [options] <offer_id1> [offer_id2] ...

Deploys transcript-variant Gemini workers to Vast.ai.
Workers are CPU-only (Gemini API calls, no GPU needed).
Gemini keys are assigned round-robin across workers (4 keys).

Options:
  --image <image>              Docker image (default: variant-worker:v2)
  --disk <gb>                  Disk in GB (default: 20)
  --worker-prefix <prefix>     Worker ID prefix (default: variant)
  --max-jobs <n>               Per-worker max shards (0 = unlimited)
  --batch-size <n>             Items per Gemini request (default: 10)
  --concurrent <n>             Concurrent Gemini requests (default: 20)
  --pack-rows <n>              Rows per output pack (default: 5000)
  --pack-videos <n>            Videos per output pack (default: 50)
  --dry-run                    Print commands without creating instances

Required env vars (loaded from .env):
  VASTAI_KEY
  R2_ENDPOINT_URL, R2_ACCESS_KEY_ID, R2_SECRET_ACCESS_KEY, R2_BUCKET
  DATABASE_URL
  GEMINI_KEY, GEMINI_PROJECT2, GEMINI_PROJECT3, GEMINI_PROJECT4
  DOCKER_USERNAME, DOCKER_PAT

Example:
  # Deploy 1 worker for smoke test
  ./scripts/deploy_variant.sh --max-jobs 2 12345678

  # Deploy 5 workers (unlimited jobs)
  ./scripts/deploy_variant.sh 11111111 22222222 33333333 44444444 55555555

  # Dry run
  ./scripts/deploy_variant.sh --dry-run 12345678 23456789
EOF
}

require_env() {
    local name="$1"
    if [ -z "${!name:-}" ]; then
        echo "ERROR: Missing required env var: ${name}" >&2
        exit 1
    fi
}

OFFERS=()
while [ $# -gt 0 ]; do
    case "$1" in
        --image)           IMAGE="$2";                          shift 2 ;;
        --disk)            DISK="$2";                           shift 2 ;;
        --worker-prefix)   WORKER_PREFIX="$2";                  shift 2 ;;
        --max-jobs)        VARIANT_MAX_JOBS="$2";               shift 2 ;;
        --batch-size)      VARIANT_BATCH_SIZE="$2";             shift 2 ;;
        --concurrent)      VARIANT_CONCURRENT_REQUESTS="$2";    shift 2 ;;
        --pack-rows)       VARIANT_PACK_TARGET_ROWS="$2";       shift 2 ;;
        --pack-videos)     VARIANT_PACK_TARGET_VIDEOS="$2";     shift 2 ;;
        --dry-run)         DRY_RUN="true";                      shift   ;;
        -h|--help)         usage; exit 0 ;;
        *)                 OFFERS+=("$1");                      shift   ;;
    esac
done

if [ ${#OFFERS[@]} -eq 0 ]; then
    usage
    exit 1
fi

require_env VASTAI_KEY
require_env R2_ENDPOINT_URL
require_env R2_ACCESS_KEY_ID
require_env R2_SECRET_ACCESS_KEY
require_env R2_BUCKET
require_env DATABASE_URL
require_env GEMINI_KEY
require_env GEMINI_PROJECT2
require_env GEMINI_PROJECT3
require_env GEMINI_PROJECT4
require_env DOCKER_USERNAME
require_env DOCKER_PAT

# 4 Gemini keys, round-robin assignment
GEMINI_KEYS=("${GEMINI_KEY}" "${GEMINI_PROJECT2}" "${GEMINI_PROJECT3}" "${GEMINI_PROJECT4}")
NUM_KEYS=${#GEMINI_KEYS[@]}

TOTAL=${#OFFERS[@]}
COUNT=0

for OFFER_ID in "${OFFERS[@]}"; do
    COUNT=$((COUNT + 1))
    KEY_INDEX=$(( (COUNT - 1) % NUM_KEYS ))
    WORKER_ID="${WORKER_PREFIX}-${OFFER_ID}"
    LABEL="${WORKER_ID}"

    # All 4 keys are passed; GEMINI_KEY_INDEX selects primary, rest are fallbacks
    ENV_STR="-e R2_ENDPOINT_URL=${R2_ENDPOINT_URL} \
-e R2_ACCESS_KEY_ID=${R2_ACCESS_KEY_ID} \
-e R2_SECRET_ACCESS_KEY=${R2_SECRET_ACCESS_KEY} \
-e R2_BUCKET=${R2_BUCKET} \
-e DATABASE_URL=${DATABASE_URL} \
-e GEMINI_KEY=${GEMINI_KEY} \
-e GEMINI_PROJECT2=${GEMINI_PROJECT2} \
-e GEMINI_PROJECT3=${GEMINI_PROJECT3} \
-e GEMINI_PROJECT4=${GEMINI_PROJECT4} \
-e GEMINI_KEY_INDEX=${KEY_INDEX} \
-e WORKER_ID=${WORKER_ID} \
-e TEMPERATURE=0 \
-e THINKING_LEVEL=low \
-e VARIANT_MAX_JOBS=${VARIANT_MAX_JOBS} \
-e VARIANT_BATCH_SIZE=${VARIANT_BATCH_SIZE} \
-e VARIANT_CONCURRENT_REQUESTS=${VARIANT_CONCURRENT_REQUESTS} \
-e VARIANT_PACK_TARGET_VIDEOS=${VARIANT_PACK_TARGET_VIDEOS} \
-e VARIANT_PACK_TARGET_ROWS=${VARIANT_PACK_TARGET_ROWS}"

    # Restart loop: if worker exits (shard error, OOM), wait 10s and restart
    ONSTART="mkdir -p /var/log/portal; printenv >> /etc/environment; cd /app; while true; do python -m src.variant_main 2>&1 | tee -a /var/log/portal/variant_worker.log; echo 'Worker exited, restarting in 10s...'; sleep 10; done &"

    echo "[${COUNT}/${TOTAL}] Creating variant worker ${WORKER_ID} (key_index=${KEY_INDEX}) from offer ${OFFER_ID}"

    CMD=(
        vastai --api-key "${VASTAI_KEY}" create instance "${OFFER_ID}"
        --image "${IMAGE}"
        --disk "${DISK}"
        --ssh --direct
        --env "${ENV_STR}"
        --login "-u ${DOCKER_USERNAME} -p ${DOCKER_PAT} docker.io"
        --onstart-cmd "${ONSTART}"
        --label "${LABEL}"
    )

    if [ "${DRY_RUN}" = "true" ]; then
        printf 'DRY RUN:'
        printf ' %q' "${CMD[@]}"
        printf '\n\n'
        continue
    fi

    "${CMD[@]}"
    echo ""
done

echo ""
if [ "${DRY_RUN}" = "true" ]; then
    echo "Dry run complete. ${TOTAL} workers previewed."
else
    echo "Deployment complete. ${TOTAL} variant workers created."
    echo ""
    echo "Monitor with:"
    echo "  vastai --api-key \"\$VASTAI_KEY\" show instances"
    echo "  vastai --api-key \"\$VASTAI_KEY\" logs <instance_id>"
    echo ""
    echo "DB monitoring:"
    echo "  SELECT worker_id, status, jobs_completed, rows_processed, active_rpm, last_heartbeat_at"
    echo "    FROM transcript_variant_workers ORDER BY last_heartbeat_at DESC;"
    echo ""
    echo "  SELECT status, count(*), sum(rows_processed) FROM transcript_variant_job_queue GROUP BY status;"
fi