#!/bin/bash
# RunPod auto-boot script
cd /workspace/

# Wait for network
echo "=== Waiting for network ==="
for i in $(seq 1 30); do
    curl -s https://pypi.org/ > /dev/null 2>&1 && echo "Network ready!" && break
    echo "Waiting... ($i)"
    sleep 2
done

# Auto-update Pod ID in .env after migration
if [ -n "$RUNPOD_POD_ID" ] && [ -f .env ]; then
    OLD_ID=$(grep '^RUNPOD_POD_ID=' .env | cut -d'=' -f2)
    if [ -n "$OLD_ID" ] && [ "$OLD_ID" != "$RUNPOD_POD_ID" ]; then
        echo "=== Pod ID changed: $OLD_ID -> $RUNPOD_POD_ID ==="
        sed -i "s/^RUNPOD_POD_ID=.*/RUNPOD_POD_ID=$RUNPOD_POD_ID/" .env
    fi
fi

if [ -f .env ]; then export $(grep -v '^#' .env | xargs); fi
PORT="${PORT:-7860}"
WHL_DIR="/workspace/.pip_cache"
TUNNEL_LOG="/tmp/cloudflared.log"

echo "=== Installing llama-cpp-python ==="
mkdir -p "$WHL_DIR"
if ls "$WHL_DIR"/llama_cpp_python-*.whl 1>/dev/null 2>&1; then
    pip install "$WHL_DIR"/llama_cpp_python-*.whl -q 2>/dev/null || true
else
    CMAKE_ARGS="-DGGML_CUDA=on" pip wheel llama-cpp-python --no-cache-dir --wheel-dir="$WHL_DIR" || true
    pip install "$WHL_DIR"/llama_cpp_python-*.whl -q 2>/dev/null || true
fi

echo "=== Upgrading torch for diffusers compatibility ==="
pip install torch==2.6.0 torchvision==0.21.0 torchaudio==2.6.0 --index-url https://download.pytorch.org/whl/cu124 -q 2>/dev/null || true

echo "=== Installing other dependencies ==="
pip install "fastapi>=0.100.0" "uvicorn>=0.20.0" "huggingface_hub>=0.23.0" "pyjwt>=2.0.0" "python-multipart>=0.0.18" -q 2>/dev/null || true
pip install diffusers==0.37.1 peft accelerate transformers sentencepiece -q 2>/dev/null || true

if [ ! -f /usr/local/bin/cloudflared ]; then
    echo "=== Installing cloudflared ==="
    wget -q https://github.com/cloudflare/cloudflared/releases/latest/download/cloudflared-linux-amd64 -O /usr/local/bin/cloudflared
    chmod +x /usr/local/bin/cloudflared
fi

pkill -f "python /workspace/server.py" 2>/dev/null || true
pkill -f cloudflared 2>/dev/null || true
sleep 1

# Clear old tunnel URL on PHP side
if [ -n "$CALLBACK_URL" ] && [ -n "$CALLBACK_SECRET" ]; then
    echo "=== Clearing old tunnel URL ==="
    curl -s -X POST "$CALLBACK_URL" -H "Content-Type: application/json" -d "{\"url\": \"clear\", \"secret\": \"$CALLBACK_SECRET\"}" || true
fi

echo "=== Starting server ==="
python -m uvicorn server:app --host 0.0.0.0 --port $PORT &
SERVER_PID=$!

echo "Waiting for server..."
for i in $(seq 1 120); do
    curl -s http://localhost:$PORT/ > /dev/null 2>&1 && echo "Server is ready!" && break
    sleep 2
done

start_tunnel() {
    pkill -f cloudflared 2>/dev/null || true
    sleep 1
    rm -f "$TUNNEL_LOG"
    cloudflared tunnel --url http://localhost:$PORT --proxy-keepalive-timeout 0s --proxy-connect-timeout 0s > "$TUNNEL_LOG" 2>&1 &
    echo "Waiting for tunnel URL..."
    for i in $(seq 1 60); do
        TUNNEL_URL=$(grep -oP 'https://[a-z0-9-]+\.trycloudflare\.com' "$TUNNEL_LOG" 2>/dev/null | head -1)
        if [ -n "$TUNNEL_URL" ]; then
            echo "=== Tunnel URL: $TUNNEL_URL ==="
            if [ -n "$CALLBACK_URL" ] && [ -n "$CALLBACK_SECRET" ]; then
                curl -s -X POST "$CALLBACK_URL" -H "Content-Type: application/json" -d "{\"url\": \"$TUNNEL_URL\", \"secret\": \"$CALLBACK_SECRET\"}"
                echo ""
                echo "=== URL sent! ==="
            fi
            return 0
        fi
        sleep 2
    done
    return 1
}

echo "=== Starting cloudflared ==="
start_tunnel

echo "=== Monitoring ==="
while kill -0 $SERVER_PID 2>/dev/null; do
    if ! pgrep -f cloudflared > /dev/null 2>&1; then
        echo "=== cloudflared died, restarting... ==="
        start_tunnel
    fi
    sleep 30
done