#!/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}/final-export-worker:latest"
DISK="${DISK:-80}"
WORKER_PREFIX="${WORKER_PREFIX:-final-export}"
STAGE="${STAGE:-video}"
DRY_RUN="false"

FINAL_EXPORT_RUN_ID="${FINAL_EXPORT_RUN_ID:-production-20260312}"
FINAL_EXPORT_MAX_VIDEOS="${FINAL_EXPORT_MAX_VIDEOS:-0}"
FINAL_EXPORT_MAX_SHARDS="${FINAL_EXPORT_MAX_SHARDS:-0}"
FINAL_EXPORT_REFERENCE_MODE="${FINAL_EXPORT_REFERENCE_MODE:-r2}"
FINAL_EXPORT_REFERENCE_BUCKET="${FINAL_EXPORT_REFERENCE_BUCKET:-${R2_BUCKET:-1-cleaned-data}}"
FINAL_EXPORT_REFERENCE_PREFIX="${FINAL_EXPORT_REFERENCE_PREFIX:-final-export-reference/20260312-v1}"
FINAL_EXPORT_OUTPUT_BUCKET="${FINAL_EXPORT_OUTPUT_BUCKET:-finalsftdata}"
FINAL_EXPORT_OUTPUT_PREFIX="${FINAL_EXPORT_OUTPUT_PREFIX:-final-export/production}"
FINAL_EXPORT_POLISH_THREADS="${FINAL_EXPORT_POLISH_THREADS:-16}"
FINAL_EXPORT_DUCKDB_THREADS="${FINAL_EXPORT_DUCKDB_THREADS:-8}"
FINAL_EXPORT_COMPACTOR_CLAIM_LIMIT="${FINAL_EXPORT_COMPACTOR_CLAIM_LIMIT:-250}"
FINAL_EXPORT_FINAL_SHARD_ROWS="${FINAL_EXPORT_FINAL_SHARD_ROWS:-15000}"
FINAL_EXPORT_MICROSHARD_TARGET_ROWS="${FINAL_EXPORT_MICROSHARD_TARGET_ROWS:-5000}"
FINAL_EXPORT_ALLOW_PARTIAL_SHARDS="${FINAL_EXPORT_ALLOW_PARTIAL_SHARDS:-false}"

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

Deploys final export workers to Vast.ai.

Options:
  --stage <video|compact>        Worker stage (default: video)
  --image <image>                Docker image
  --disk <gb>                    Disk size in GB
  --worker-prefix <prefix>       Worker prefix
  --max-videos <n>               Max videos for video-stage worker
  --max-shards <n>               Max shards for compactor worker
  --output-prefix <prefix>       Final export output prefix
  --reference-prefix <prefix>    Reference snapshot prefix
  --allow-partial-shards         Allow partial shards for compact stage
  --dry-run                      Print commands only
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
        --stage)                  STAGE="$2"; shift 2 ;;
        --image)                  IMAGE="$2"; shift 2 ;;
        --disk)                   DISK="$2"; shift 2 ;;
        --worker-prefix)          WORKER_PREFIX="$2"; shift 2 ;;
        --max-videos)             FINAL_EXPORT_MAX_VIDEOS="$2"; shift 2 ;;
        --max-shards)             FINAL_EXPORT_MAX_SHARDS="$2"; shift 2 ;;
        --output-prefix)          FINAL_EXPORT_OUTPUT_PREFIX="$2"; shift 2 ;;
        --reference-prefix)       FINAL_EXPORT_REFERENCE_PREFIX="$2"; shift 2 ;;
        --allow-partial-shards)   FINAL_EXPORT_ALLOW_PARTIAL_SHARDS="true"; shift ;;
        --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 DATABASE_URL
require_env DOCKER_USERNAME
require_env DOCKER_PAT

COUNT=0
TOTAL=${#OFFERS[@]}
for OFFER_ID in "${OFFERS[@]}"; do
    COUNT=$((COUNT + 1))
    WORKER_ID="${WORKER_PREFIX}-${STAGE}-${OFFER_ID}"
    LABEL="${WORKER_ID}"

    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 DATABASE_URL=${DATABASE_URL} \
-e WORKER_ID=${WORKER_ID} \
-e FINAL_EXPORT_STAGE=${STAGE} \
-e FINAL_EXPORT_REFERENCE_MODE=${FINAL_EXPORT_REFERENCE_MODE} \
-e FINAL_EXPORT_REFERENCE_BUCKET=${FINAL_EXPORT_REFERENCE_BUCKET} \
-e FINAL_EXPORT_REFERENCE_PREFIX=${FINAL_EXPORT_REFERENCE_PREFIX} \
-e FINAL_EXPORT_OUTPUT_BUCKET=${FINAL_EXPORT_OUTPUT_BUCKET} \
-e FINAL_EXPORT_OUTPUT_PREFIX=${FINAL_EXPORT_OUTPUT_PREFIX} \
-e FINAL_EXPORT_POLISH_THREADS=${FINAL_EXPORT_POLISH_THREADS} \
-e FINAL_EXPORT_DUCKDB_THREADS=${FINAL_EXPORT_DUCKDB_THREADS} \
-e FINAL_EXPORT_COMPACTOR_CLAIM_LIMIT=${FINAL_EXPORT_COMPACTOR_CLAIM_LIMIT} \
-e FINAL_EXPORT_FINAL_SHARD_ROWS=${FINAL_EXPORT_FINAL_SHARD_ROWS} \
-e FINAL_EXPORT_MICROSHARD_TARGET_ROWS=${FINAL_EXPORT_MICROSHARD_TARGET_ROWS} \
-e FINAL_EXPORT_ALLOW_PARTIAL_SHARDS=${FINAL_EXPORT_ALLOW_PARTIAL_SHARDS} \
-e FINAL_EXPORT_RUN_ID=${FINAL_EXPORT_RUN_ID} \
-e FINAL_EXPORT_MAX_VIDEOS=${FINAL_EXPORT_MAX_VIDEOS} \
-e FINAL_EXPORT_MAX_SHARDS=${FINAL_EXPORT_MAX_SHARDS}"

    ONSTART="mkdir -p /var/log/portal; printenv >> /etc/environment; cd /app; while true; do /app/scripts/run_final_export_worker.sh 2>&1 | tee -a /var/log/portal/final_export_${STAGE}.log; echo 'Worker exited, restarting in 10s...'; sleep 10; done &"

    echo "[${COUNT}/${TOTAL}] Creating ${STAGE} final export worker ${WORKER_ID} 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

if [ "${DRY_RUN}" = "true" ]; then
    echo "Dry run complete. ${TOTAL} final export workers previewed."
else
    echo "Deployment complete. ${TOTAL} final export workers created."
fi