#!/bin/bash
set -x
WHEEL_FILE="/workspace/rl4phyx/flash_attn-2.8.3+cu12torch2.6cxx11abiFALSE-cp311-cp311-linux_x86_64.whl"
rm -f "$WHEEL_FILE"

# Try each mirror
MIRRORS=(
    "https://ghfast.top/https://github.com/Dao-AILab/flash-attention/releases/download/v2.8.3/flash_attn-2.8.3%2Bcu12torch2.6cxx11abiFALSE-cp311-cp311-linux_x86_64.whl"
    "https://github.moeyy.xyz/https://github.com/Dao-AILab/flash-attention/releases/download/v2.8.3/flash_attn-2.8.3%2Bcu12torch2.6cxx11abiFALSE-cp311-cp311-linux_x86_64.whl"
    "https://ghps.cc/https://github.com/Dao-AILab/flash-attention/releases/download/v2.8.3/flash_attn-2.8.3%2Bcu12torch2.6cxx11abiFALSE-cp311-cp311-linux_x86_64.whl"
    "https://gh-proxy.com/https://github.com/Dao-AILab/flash-attention/releases/download/v2.8.3/flash_attn-2.8.3%2Bcu12torch2.6cxx11abiFALSE-cp311-cp311-linux_x86_64.whl"
    "https://mirror.ghproxy.com/https://github.com/Dao-AILab/flash-attention/releases/download/v2.8.3/flash_attn-2.8.3%2Bcu12torch2.6cxx11abiFALSE-cp311-cp311-linux_x86_64.whl"
    "https://ghproxy.net/https://github.com/Dao-AILab/flash-attention/releases/download/v2.8.3/flash_attn-2.8.3%2Bcu12torch2.6cxx11abiFALSE-cp311-cp311-linux_x86_64.whl"
    "https://gh.api.99988866.xyz/https://github.com/Dao-AILab/flash-attention/releases/download/v2.8.3/flash_attn-2.8.3%2Bcu12torch2.6cxx11abiFALSE-cp311-cp311-linux_x86_64.whl"
    "https://github.com/Dao-AILab/flash-attention/releases/download/v2.8.3/flash_attn-2.8.3%2Bcu12torch2.6cxx11abiFALSE-cp311-cp311-linux_x86_64.whl"
)

for i in "${!MIRRORS[@]}"; do
    URL="${MIRRORS[$i]}"
    echo "=== Mirror $((i+1))/${#MIRRORS[@]}: ${URL:0:80}... ==="
    wget -q --timeout=30 --tries=2 -O "$WHEEL_FILE" "$URL" 2>&1
    if [ -f "$WHEEL_FILE" ] && [ $(stat -c%s "$WHEEL_FILE" 2>/dev/null || echo 0) -gt 1000000 ]; then
        echo "SUCCESS! Downloaded $(ls -lh $WHEEL_FILE | awk '{print $5}')"
        break
    else
        echo "Failed, trying next..."
        rm -f "$WHEEL_FILE"
    fi
done

# Also try pip install directly with mirror URL
if [ ! -f "$WHEEL_FILE" ] || [ $(stat -c%s "$WHEEL_FILE" 2>/dev/null || echo 0) -lt 1000000 ]; then
    echo "=== Trying pip install directly ==="
    for URL in "${MIRRORS[@]}"; do
        echo "pip install ${URL:0:80}..."
        pip install "$URL" --no-deps 2>&1 | tail -5
        python3 -c "import flash_attn; print('OK')" 2>/dev/null && echo "PIP_INSTALL_OK" && break
    done
fi

# If we have the wheel, install it
if [ -f "$WHEEL_FILE" ] && [ $(stat -c%s "$WHEEL_FILE" 2>/dev/null || echo 0) -gt 1000000 ]; then
    echo "=== Installing from wheel ==="
    pip install "$WHEEL_FILE" --force-reinstall --no-deps 2>&1 | tail -5
fi

# Verify
echo "=== Verify ==="
python3 -c "import flash_attn; print(f'flash_attn: {flash_attn.__version__}')" 2>&1
echo "MIRROR_SCRIPT_DONE"