#!/bin/bash
# Deploy Azaion detections demo stack to a Jetson over SSH.
#
# Usage:
#   JETSON_HOST=192.168.x.x bash scripts/deploy_demo_jetson.sh \
#     --onnx      /local/path/azaion.onnx \
#     --classes   /local/path/classes.json \
#     [--int8-cache /local/path/azaion.int8_calib.cache] \
#     [--calibration-images /local/path/images/]
#
# Optional env vars:
#   JETSON_HOST   (required) IP or hostname of the Jetson
#   JETSON_USER   SSH user (default: jetson)
#   REMOTE_DIR    Path on Jetson to deploy into (default: ~/detections)
set -euo pipefail

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
PROJECT_ROOT="$(dirname "$SCRIPT_DIR")"

JETSON_HOST="${JETSON_HOST:-}"
JETSON_USER="${JETSON_USER:-jetson}"
REMOTE_DIR="${REMOTE_DIR:-~/detections}"

ONNX_PATH=""
CLASSES_PATH=""
INT8_CACHE_PATH=""
CALIBRATION_IMAGES=""

usage() {
    echo "Usage: JETSON_HOST=<ip> bash $0 --onnx <path> --classes <path> [options]"
    echo ""
    echo "Required:"
    echo "  --onnx <path>            Local path to azaion.onnx"
    echo "  --classes <path>         Local path to classes.json"
    echo ""
    echo "Optional:"
    echo "  --int8-cache <path>      Local path to azaion.int8_calib.cache (skips calibration)"
    echo "  --calibration-images <dir>  Local image directory; rsync to Jetson and run INT8 calibration"
    echo "  --help                   Show this message"
    echo ""
    echo "Env vars:"
    echo "  JETSON_HOST    (required) Jetson IP or hostname"
    echo "  JETSON_USER    SSH user (default: jetson)"
    echo "  REMOTE_DIR     Deploy directory on Jetson (default: ~/detections)"
    exit 0
}

while [[ $# -gt 0 ]]; do
    case "$1" in
        --onnx)              ONNX_PATH="$2";             shift 2 ;;
        --classes)           CLASSES_PATH="$2";          shift 2 ;;
        --int8-cache)        INT8_CACHE_PATH="$2";       shift 2 ;;
        --calibration-images) CALIBRATION_IMAGES="$2";  shift 2 ;;
        --help)              usage ;;
        *) echo "Unknown argument: $1"; usage ;;
    esac
done

[[ -z "$JETSON_HOST" ]]  && { echo "ERROR: JETSON_HOST is required"; exit 1; }
[[ -z "$ONNX_PATH" ]]    && { echo "ERROR: --onnx is required"; exit 1; }
[[ -z "$CLASSES_PATH" ]] && { echo "ERROR: --classes is required"; exit 1; }
[[ -f "$ONNX_PATH" ]]    || { echo "ERROR: ONNX file not found: $ONNX_PATH"; exit 1; }
[[ -f "$CLASSES_PATH" ]] || { echo "ERROR: classes.json not found: $CLASSES_PATH"; exit 1; }

SSH="ssh ${JETSON_USER}@${JETSON_HOST}"
SCP="scp"

echo "=== Azaion Demo — Jetson Deployment ==="
echo "  Host:       ${JETSON_USER}@${JETSON_HOST}"
echo "  Remote dir: ${REMOTE_DIR}"
echo ""

# ── 1. Sync project ─────────────────────────────────────────────────────────
echo "--- Syncing project files ---"
$SSH "mkdir -p ${REMOTE_DIR}/demo/models"
rsync -az --exclude='.git' --exclude='__pycache__' --exclude='*.pyc' \
    --exclude='*.egg-info' --exclude='.venv' --exclude='demo/models' \
    "${PROJECT_ROOT}/" "${JETSON_USER}@${JETSON_HOST}:${REMOTE_DIR}/"

# ── 2. Upload model artifacts ────────────────────────────────────────────────
echo "--- Uploading model artifacts ---"
$SCP "$ONNX_PATH"    "${JETSON_USER}@${JETSON_HOST}:${REMOTE_DIR}/demo/models/azaion.onnx"
$SCP "$CLASSES_PATH" "${JETSON_USER}@${JETSON_HOST}:${REMOTE_DIR}/demo/models/classes.json"

if [[ -n "$INT8_CACHE_PATH" ]]; then
    [[ -f "$INT8_CACHE_PATH" ]] || { echo "ERROR: INT8 cache not found: $INT8_CACHE_PATH"; exit 1; }
    echo "--- Uploading INT8 calibration cache ---"
    $SCP "$INT8_CACHE_PATH" "${JETSON_USER}@${JETSON_HOST}:${REMOTE_DIR}/demo/models/azaion.int8_calib.cache"
fi

# ── 3. Optional: run INT8 calibration on Jetson ──────────────────────────────
if [[ -n "$CALIBRATION_IMAGES" ]] && [[ -z "$INT8_CACHE_PATH" ]]; then
    [[ -d "$CALIBRATION_IMAGES" ]] || { echo "ERROR: calibration images dir not found: $CALIBRATION_IMAGES"; exit 1; }
    echo "--- Syncing calibration images to Jetson ---"
    $SSH "mkdir -p ${REMOTE_DIR}/demo/calibration"
    rsync -az "${CALIBRATION_IMAGES}/" "${JETSON_USER}@${JETSON_HOST}:${REMOTE_DIR}/demo/calibration/"

    echo "--- Building detections image for calibration ---"
    $SSH "cd ${REMOTE_DIR} && docker compose -f docker-compose.demo-jetson.yml build detections"

    echo "--- Running INT8 calibration (this takes several minutes) ---"
    $SSH "cd ${REMOTE_DIR} && docker compose -f docker-compose.demo-jetson.yml run --rm \
        -v ${REMOTE_DIR}/demo/calibration:/calibration \
        detections \
        python3 scripts/generate_int8_cache.py \
            --images-dir /calibration \
            --onnx /models/azaion.onnx \
            --output /models/azaion.int8_calib.cache"
    echo "--- Calibration cache written to ${REMOTE_DIR}/demo/models/azaion.int8_calib.cache ---"
fi

# ── 4. Start services ────────────────────────────────────────────────────────
echo "--- Building and starting services ---"
$SSH "cd ${REMOTE_DIR} && docker compose -f docker-compose.demo-jetson.yml up -d --build"

# ── 5. Health check ───────────────────────────────────────────────────────────
echo "--- Health check ---"
HEALTH_URL="http://${JETSON_HOST}:8080/health"
MAX_RETRIES=15
RETRY_INTERVAL=5

for i in $(seq 1 $MAX_RETRIES); do
    STATUS=$(curl -s -o /dev/null -w "%{http_code}" "$HEALTH_URL" 2>/dev/null || true)
    if [[ "$STATUS" == "200" ]]; then
        echo ""
        echo "=== Demo is live at http://${JETSON_HOST}:8080 ==="
        echo ""
        echo "Endpoints:"
        echo "  POST http://${JETSON_HOST}:8080/detect/image"
        echo "  POST http://${JETSON_HOST}:8080/detect/video"
        echo "  GET  http://${JETSON_HOST}:8080/health"
        echo ""
        echo "Note: On first start the service converts azaion.onnx to a TRT engine."
        echo "      Check /health until AI status shows 'enabled'."
        exit 0
    fi
    echo "  Waiting for service… (${i}/${MAX_RETRIES}, HTTP ${STATUS})"
    sleep "$RETRY_INTERVAL"
done

echo "ERROR: Health check failed after $((MAX_RETRIES * RETRY_INTERVAL))s"
echo "Check logs with: ssh ${JETSON_USER}@${JETSON_HOST} \"cd ${REMOTE_DIR} && docker compose -f docker-compose.demo-jetson.yml logs detections\""
exit 1