#!/bin/bash
set -e

REPO_URL="https://git.xiaoxiajianji.com/xiaoxia/xiaoxia-saas/raw/branch/develop/deploy/gpu_worker"
MUSE_DIR="/home/ying/projects/MuseTalk"
WORKER_DIR="/opt/xiaoxia-gpu-worker"
LOG_FILE="$HOME/gpu-worker-update.log"

log() {
    local NOW
    NOW=$(date +"%Y-%m-%d %H:%M:%S")
    echo "[$NOW] $*" | tee -a "$LOG_FILE"
}

log "========== start update =========="

BAK_SUFFIX=$(date +"%Y%m%d%H%M%S")
cp "$MUSE_DIR/musetalk_server.py" "$MUSE_DIR/musetalk_server.py.bak.$BAK_SUFFIX"
cp "$WORKER_DIR/gpu_worker.py" "$WORKER_DIR/gpu_worker.py.bak.$BAK_SUFFIX"
log "backup done ($BAK_SUFFIX)"

wget -q -O "$MUSE_DIR/musetalk_server.py" "$REPO_URL/musetalk_server.py"
log "musetalk_server.py updated"

wget -q -O "$WORKER_DIR/gpu_worker.py" "$REPO_URL/gpu_worker.py"
log "gpu_worker.py updated"

wget -q -O /tmp/gpu-requirements.txt "$REPO_URL/requirements.txt"
if [ -f "$WORKER_DIR/requirements.txt" ] && ! diff -q "$WORKER_DIR/requirements.txt" /tmp/gpu-requirements.txt > /dev/null 2>&1; then
    log "requirements changed, updating..."
    cp /tmp/gpu-requirements.txt "$WORKER_DIR/requirements.txt"
    "$WORKER_DIR/venv/bin/pip" install -r "$WORKER_DIR/requirements.txt" -q
    log "pip install done"
else
    log "requirements no change, skip pip"
fi

sudo systemctl restart musetalk-worker
log "musetalk restarted"
sleep 5

sudo systemctl restart xiaoxia-gpu-worker
log "gpu-worker restarted"
sleep 8

HEALTH=$(curl -s http://127.0.0.1:7861/health 2>/dev/null)
if echo "$HEALTH" | grep -q "healthy\|ok"; then
    log "health check OK"
    log "========== update done =========="
    exit 0
else
    log "health check FAILED, rolling back..."
    LATEST_MUSE_BAK=$(ls -t "$MUSE_DIR/musetalk_server.py.bak."* 2>/dev/null | head -1)
    LATEST_WORKER_BAK=$(ls -t "$WORKER_DIR/gpu_worker.py.bak."* 2>/dev/null | head -1)
    [ -n "$LATEST_MUSE_BAK" ] && cp "$LATEST_MUSE_BAK" "$MUSE_DIR/musetalk_server.py"
    [ -n "$LATEST_WORKER_BAK" ] && cp "$LATEST_WORKER_BAK" "$WORKER_DIR/gpu_worker.py"
    sudo systemctl restart musetalk-worker
    sleep 5
    sudo systemctl restart xiaoxia-gpu-worker
    log "rolled back"
    exit 1
fi