#!/bin/bash
set -x

echo "=== Step 1: Install veRL from GitHub main ==="
# Uninstall old version
pip uninstall verl -y 2>&1 | tail -3

# Install from GitHub via pip
# Try ghfast mirror first (for China), fallback to direct
pip install 'git+https://ghfast.top/https://github.com/volcengine/verl.git@main' --no-deps \
    2>&1 | tail -10

# If failed, try direct GitHub
if ! python3 -c "import verl" 2>/dev/null; then
    echo "Mirror failed, trying direct GitHub..."
    pip install 'git+https://github.com/volcengine/verl.git@main' --no-deps \
        2>&1 | tail -10
fi

echo ""
echo "=== Step 2: Clear pycache ==="
VERL_PATH=$(python3 -c "import verl; import os; print(os.path.dirname(verl.__file__))")
echo "veRL path: $VERL_PATH"
find "$VERL_PATH" -type d -name __pycache__ -exec rm -rf {} + 2>/dev/null

echo ""
echo "=== Step 3: Patch flash_attention_2 -> sdpa ==="
find "$VERL_PATH" -name "*.py" -exec sed -i "s/flash_attention_2/sdpa/g" {} +
find "$VERL_PATH" -name "*.yaml" -exec sed -i "s/flash_attention_2/sdpa/g" {} +
echo "Remaining flash_attention_2: $(grep -rn 'flash_attention_2' $VERL_PATH --include='*.py' 2>/dev/null | wc -l)"
echo "sdpa refs: $(grep -rn '"sdpa"' $VERL_PATH --include='*.py' --include='*.yaml' 2>/dev/null | wc -l)"
find "$VERL_PATH" -type d -name __pycache__ -exec rm -rf {} + 2>/dev/null

echo ""
echo "=== Step 4: Verify ==="
python3 << 'PYEOF'
import torch
print(f"torch: {torch.__version__}")

import vllm
print(f"vLLM: {vllm.__version__}")

import verl
print(f"veRL: OK (path: {verl.__file__})")

import transformers
print(f"transformers: {transformers.__version__}")

# Check if vllm_async_server can import properly now
try:
    from verl.workers.rollout.vllm_rollout.vllm_async_server import vLLMHttpServer
    print("vLLMHttpServer: OK")
except ImportError as e:
    print(f"vLLMHttpServer import error: {e}")
    # If it fails, check what it needs
    import traceback
    traceback.print_exc()

print("VERIFY_DONE")
PYEOF

echo "INSTALL_DONE"