#!/bin/bash
#
# RunPod Initialization Script
# Automatically downloads snapshot from HuggingFace, extracts to RAM disk, and starts Service Manager
#
# Usage: Used as --args in runpodctl create pod command
#

set -e  # Exit on error

echo "════════════════════════════════════════════════════════════════"
echo "🚀 Ultravox Pipeline - RunPod Initialization"
echo "════════════════════════════════════════════════════════════════"
echo ""

# Configuration
SNAPSHOT_URL="${SNAPSHOT_URL:-https://huggingface.co/marcosremar2/ultravox-snapshots/resolve/main/ultravox-snapshot-latest.tar.zst}"
PROJECT_DIR="/root/ultravox-pipeline"
RAMDISK_PATH="/dev/shm"
ULTRAVOX_PROFILE="${ULTRAVOX_PROFILE:-gpu-prod}"

# Detect compression format from URL
if [[ "$SNAPSHOT_URL" == *.tar.zst ]]; then
    COMPRESS_FORMAT="zstd"
    FILE_EXT=".tar.zst"
    TAR_FLAGS="--zstd -xf"
elif [[ "$SNAPSHOT_URL" == *.tar.gz ]] || [[ "$SNAPSHOT_URL" == *.tgz ]]; then
    COMPRESS_FORMAT="gzip"
    FILE_EXT=".tar.gz"
    TAR_FLAGS="-xzf"
else
    COMPRESS_FORMAT="unknown"
    FILE_EXT=".tar.gz"
    TAR_FLAGS="-xzf"
    echo "⚠️  Warning: Unknown compression format, assuming gzip"
fi

TEMP_SNAPSHOT="$RAMDISK_PATH/snapshot$FILE_EXT"

echo "📋 Configuration:"
echo "   Snapshot URL: $SNAPSHOT_URL"
echo "   Compression:  $COMPRESS_FORMAT"
echo "   Project Dir:  $PROJECT_DIR"
echo "   RAM Disk:     $RAMDISK_PATH"
echo "   Profile:      $ULTRAVOX_PROFILE"
echo ""

# Step 1: Check available RAM
echo "🔍 Step 1/5: Checking RAM disk space..."
df -h $RAMDISK_PATH
echo ""

# Step 2: Download snapshot to RAM disk (ultra-fast)
echo "📥 Step 2/5: Downloading snapshot to RAM disk..."
START_TIME=$(date +%s)

wget -q --show-progress -O "$TEMP_SNAPSHOT" "$SNAPSHOT_URL" || {
    echo "⚠️  wget failed, trying curl..."
    curl -# -L -o "$TEMP_SNAPSHOT" "$SNAPSHOT_URL"
}

DOWNLOAD_TIME=$(($(date +%s) - START_TIME))
SNAPSHOT_SIZE=$(du -sh "$TEMP_SNAPSHOT" | cut -f1)
echo "✅ Downloaded: $SNAPSHOT_SIZE in ${DOWNLOAD_TIME}s"
echo ""

# Step 3: Extract from RAM disk to project directory (ultra-fast)
echo "📦 Step 3/5: Extracting snapshot from RAM disk..."
echo "   Using tar flags: $TAR_FLAGS"
START_TIME=$(date +%s)

mkdir -p "$PROJECT_DIR"
cd "$PROJECT_DIR"
tar $TAR_FLAGS "$TEMP_SNAPSHOT"

EXTRACT_TIME=$(($(date +%s) - START_TIME))
echo "✅ Extracted in ${EXTRACT_TIME}s"
echo ""

# Step 3.5: Validate extraction
echo "🔍 Step 3.5/6: Validating extraction..."
VALIDATION_FAILED=0

if [ ! -d "$PROJECT_DIR/src" ]; then
    echo "❌ Missing: src/"
    VALIDATION_FAILED=1
fi

if [ ! -d "$PROJECT_DIR/config" ]; then
    echo "❌ Missing: config/"
    VALIDATION_FAILED=1
fi

if [ ! -f "$PROJECT_DIR/src/core/service_manager/main.py" ]; then
    echo "❌ Missing: src/core/service_manager/main.py"
    VALIDATION_FAILED=1
fi

if [ $VALIDATION_FAILED -eq 1 ]; then
    echo "❌ Extraction validation failed!"
    echo "📂 Contents of $PROJECT_DIR:"
    ls -la "$PROJECT_DIR"
    exit 1
fi

echo "✅ Validation passed - all required files present"
echo ""

# Step 4: Cleanup RAM disk
echo "🧹 Step 4/6: Cleaning up RAM disk..."
rm -f "$TEMP_SNAPSHOT"
echo "✅ RAM disk cleaned"
echo ""

# Step 5: Start Service Manager
echo "⚙️  Step 5/6: Starting Service Manager..."
cd "$PROJECT_DIR"

# Export profile and authorization
export ULTRAVOX_PROFILE="$ULTRAVOX_PROFILE"
export ULTRAVOX_SM_AUTHORIZED=1

# Check if Service Manager main.py exists
if [ ! -f "src/core/service_manager/main.py" ]; then
    echo "❌ src/core/service_manager/main.py not found in $PROJECT_DIR"
    echo "📂 Contents:"
    ls -la
    exit 1
fi

# Start Service Manager in background using Python directly
nohup python3 src/core/service_manager/main.py > /var/log/service_manager.log 2>&1 &
SERVICE_MANAGER_PID=$!

echo "✅ Service Manager started (PID: $SERVICE_MANAGER_PID)"
echo "   Logs: /var/log/service_manager.log"
echo ""

# Wait a few seconds for initialization
sleep 5

# Check if Service Manager is running
if kill -0 $SERVICE_MANAGER_PID 2>/dev/null; then
    echo "✅ Service Manager is running"
else
    echo "❌ Service Manager failed to start"
    echo "📋 Last 20 lines of log:"
    tail -n 20 /var/log/service_manager.log
    exit 1
fi

echo ""
echo "════════════════════════════════════════════════════════════════"
echo "✅ INITIALIZATION COMPLETE!"
echo "════════════════════════════════════════════════════════════════"
echo ""
echo "📊 Timing Summary:"
TOTAL_TIME=$((DOWNLOAD_TIME + EXTRACT_TIME))
echo "   Download:  ${DOWNLOAD_TIME}s"
echo "   Extract:   ${EXTRACT_TIME}s"
echo "   Total:     ${TOTAL_TIME}s"
echo ""
echo "🎯 Service Manager Status:"
echo "   PID:       $SERVICE_MANAGER_PID"
echo "   Profile:   $ULTRAVOX_PROFILE"
echo "   Logs:      /var/log/service_manager.log"
echo ""
echo "🔍 Check status:"
echo "   curl http://localhost:8888/health"
echo ""
echo "📝 View logs:"
echo "   tail -f /var/log/service_manager.log"
echo ""

# Keep container running
exec tail -f /var/log/service_manager.log