#!/bin/bash

# Define mount environments
APP_STORAGE="/app/storage"
PERSISTENT_STORAGE="/data/storage"

# ApexKit persistent file/database storage
export APEXKIT_MOUNTED_FILE_STORAGE="/data/"

# Force WASM engine to strictly use the fast local NVMe cache in /app
export APEXKIT_WASM_CACHE_DIR="/app/.cache/wasm"

echo "[Init] Setting up Hugging Face Storage Sync..."
echo "[Init] Persistent File Storage: ${APEXKIT_MOUNTED_FILE_STORAGE}"
echo "[Init] Fast NVMe WASM Cache:   ${APEXKIT_WASM_CACHE_DIR}"

# 0. Restore upgraded binary and rollbacks from persistent storage if present
if [ -d "/data/bin" ] && [ -f "/data/bin/apexkit" ]; then
    echo "[Sync] Restoring upgraded apexkit binary and rollbacks from persistent storage..."
    mkdir -p /app/bin
    cp -a /data/bin/* /app/bin/ 2>/dev/null || :
    chmod +x /app/bin/apexkit* 2>/dev/null || :
fi

# 1. Restore database data from persistent storage on boot (EXCLUDING uploads)
if [ -d "$PERSISTENT_STORAGE" ] && [ "$(ls -A $PERSISTENT_STORAGE 2>/dev/null)" ]; then
    echo "[Sync] Restoring data from persistent storage ($PERSISTENT_STORAGE) to NVMe ($APP_STORAGE)..."
    mkdir -p "$APP_STORAGE"
    rsync -a --exclude="**/uploads/" "$PERSISTENT_STORAGE/" "$APP_STORAGE/"
    echo "[Sync] Restore complete."
else
    echo "[Sync] No existing persistent data found. Starting fresh."
    mkdir -p "$APP_STORAGE"
    mkdir -p "$PERSISTENT_STORAGE"
fi

# Inject Replica ID if provided via Env Vars
if [ -n "$APEX_REPLICA_ID" ]; then
    mkdir -p "$APP_STORAGE/system"
    echo "$APEX_REPLICA_ID" > "$APP_STORAGE/system/.replica_id"
    echo "[Init] Injected static Replica ID: $APEX_REPLICA_ID"
fi

# 2. Background Sync Loop (Every 60 minutes)
sync_data() {
    while true; do
        sleep 3600
        echo "[Sync] Backing up NVMe data to persistent storage..."
        rsync -a --delete --exclude="**/uploads/" "$APP_STORAGE/" "$PERSISTENT_STORAGE/"
        if [ -d "/data" ]; then
            rsync -a --delete /app/bin/ /data/bin/ 2>/dev/null || :
        fi
        echo "[Sync] Backup complete."
    done
}

sync_data &
SYNC_PID=$!

# 3. Graceful Shutdown Handler
cleanup() {
    echo "[Shutdown] Caught stop signal. Stopping application..."
    
    pkill -TERM -x apexkit 2>/dev/null
    kill -TERM "$FASTAPI_PID" 2>/dev/null
    
    wait 2>/dev/null
    
    echo "[Shutdown] Performing final data sync to persistent storage..."
    rsync -a --delete --exclude="**/uploads/" "$APP_STORAGE/" "$PERSISTENT_STORAGE/"
    if [ -d "/data" ]; then
        rsync -a --delete /app/bin/ /data/bin/ 2>/dev/null || :
    fi
    echo "[Shutdown] Final sync complete. Safely exiting."
    
    kill -9 "$SYNC_PID" 2>/dev/null
    exit 0
}

trap cleanup SIGINT SIGTERM

# 4. Start FastAPI background server on port 8000
echo "[App] Starting FastAPI on port 8000..."
uvicorn main:app --host 127.0.0.1 --port 8000 --workers 1 &
FASTAPI_PID=$!

# 5. Start ApexKit on the exposed port
echo "[App] Starting ApexKit on port ${PORT:-7860}..."
apexkit --port ${PORT:-7860} &
APP_PID=$!

# 6. Supervise ApexKit across live zero-downtime upgrades
while [ -n "$APP_PID" ] && kill -0 "$APP_PID" 2>/dev/null; do
    # Block until $APP_PID exits (works for both child and reparented processes)
    tail --pid="$APP_PID" -f /dev/null 2>/dev/null

    # Allow the process table to settle
    sleep 2

    # Check if a new replacement process was spawned during an upgrade/rollback
    NEW_PID=$(pgrep -x apexkit | head -n 1)

    if [ -n "$NEW_PID" ] && [ "$NEW_PID" != "$APP_PID" ]; then
        echo "[App] Upgraded ApexKit process running (PID: $NEW_PID). Continuing supervision..."
        APP_PID="$NEW_PID"

        # Sync the new binary and rollbacks to persistent disk ONCE
        if [ -d "/data" ]; then
            mkdir -p /data/bin
            rsync -a --delete /app/bin/ /data/bin/ 2>/dev/null || :
        fi
        continue
    fi

    # No replacement process found; the server has shut down
    echo "[App] ApexKit has stopped."
    break
done

echo "[App] Application exited. Syncing data before container stops..."
kill -TERM "$FASTAPI_PID" 2>/dev/null
rsync -a --delete --exclude="**/uploads/" "$APP_STORAGE/" "$PERSISTENT_STORAGE/"
if [ -d "/data" ]; then
    rsync -a --delete /app/bin/ /data/bin/ 2>/dev/null || :
fi
kill -9 "$SYNC_PID" 2>/dev/null