#!/bin/bash

# Setup Suno Environment with Flash Attention v3 (Hopper/H100 optimized)
# This version can use a custom Flash Attention path or clone if missing

set -e  # Exit on error

# Configuration
ENV_NAME="${1:-suno_env_fa3}"
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
LOG_FILE="${SCRIPT_DIR}/tmp/setup_${ENV_NAME}_$(date +%Y%m%d_%H%M%S).log"

# Create tmp directory if it doesn't exist
mkdir -p "${SCRIPT_DIR}/tmp"

# Colors for output
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
NC='\033[0m' # No Color

# Logging functions
log() {
    echo -e "${GREEN}[$(date +'%Y-%m-%d %H:%M:%S')]${NC} $1" | tee -a "$LOG_FILE"
}

error() {
    echo -e "${RED}[ERROR]${NC} $1" | tee -a "$LOG_FILE"
    exit 1
}

warning() {
    echo -e "${YELLOW}[WARNING]${NC} $1" | tee -a "$LOG_FILE"
}

info() {
    echo -e "${BLUE}[INFO]${NC} $1" | tee -a "$LOG_FILE"
}

# Function to setup Suno Utils path
setup_suno_utils_path() {
    # Default path
    DEFAULT_SUNO_PATH="/home/vibert/projects/glockenspiel/suno_utils"
    
    # Ask user for Suno Utils path
    echo -e "${BLUE}Suno Utils Setup${NC}"
    echo "Enter the path to suno_utils package"
    echo "Press Enter to use default: $DEFAULT_SUNO_PATH"
    read -p "Path: " USER_SUNO_PATH
    
    # Use user path or default
    if [ -z "$USER_SUNO_PATH" ]; then
        SUNO_UTILS_PATH="$DEFAULT_SUNO_PATH"
        log "Using default suno_utils path: $SUNO_UTILS_PATH"
    else
        # Expand tilde if present
        SUNO_UTILS_PATH="${USER_SUNO_PATH/#\~/$HOME}"
        log "Using custom suno_utils path: $SUNO_UTILS_PATH"
    fi
    
    # Check if suno_utils exists
    if [ ! -d "$SUNO_UTILS_PATH" ]; then
        error "suno_utils not found at $SUNO_UTILS_PATH. Please provide a valid path to the suno_utils package."
    else
        log "suno_utils found at $SUNO_UTILS_PATH"
        # Check if it's a valid Python package
        if [ ! -f "$SUNO_UTILS_PATH/setup.py" ] && [ ! -f "$SUNO_UTILS_PATH/pyproject.toml" ]; then
            warning "$SUNO_UTILS_PATH doesn't appear to be a valid Python package"
            warning "Missing setup.py or pyproject.toml"
            read -p "Continue anyway? (y/n): " -n 1 -r
            echo
            if [[ ! $REPLY =~ ^[Yy]$ ]]; then
                error "Valid suno_utils package required"
            fi
        fi
    fi
}

# Function to setup Flash Attention path
setup_flash_attention_path() {
    # Default path
    DEFAULT_FLASH_PATH="$HOME/projects/flash-attention"
    
    # Ask user for Flash Attention path
    echo -e "${BLUE}Flash Attention v3 Setup${NC}"
    echo "Enter the path to Flash Attention repository"
    echo "Press Enter to use default: $DEFAULT_FLASH_PATH"
    read -p "Path: " USER_FLASH_PATH
    
    # Use user path or default
    if [ -z "$USER_FLASH_PATH" ]; then
        FLASH_ATTN_PATH="$DEFAULT_FLASH_PATH"
        log "Using default Flash Attention path: $FLASH_ATTN_PATH"
    else
        # Expand tilde if present
        FLASH_ATTN_PATH="${USER_FLASH_PATH/#\~/$HOME}"
        log "Using custom Flash Attention path: $FLASH_ATTN_PATH"
    fi
    
    # Check if Flash Attention exists
    if [ ! -d "$FLASH_ATTN_PATH" ]; then
        warning "Flash Attention not found at $FLASH_ATTN_PATH"
        read -p "Would you like to clone Flash Attention to this location? (y/n): " -n 1 -r
        echo
        if [[ $REPLY =~ ^[Yy]$ ]]; then
            log "Cloning Flash Attention repository..."
            # Create parent directory if needed
            mkdir -p "$(dirname "$FLASH_ATTN_PATH")"
            git clone https://github.com/Dao-AILab/flash-attention.git "$FLASH_ATTN_PATH" >> "$LOG_FILE" 2>&1 || {
                error "Failed to clone Flash Attention repository"
            }
            
            # Initialize submodules
            cd "$FLASH_ATTN_PATH"
            log "Initializing Flash Attention submodules..."
            git submodule update --init --recursive >> "$LOG_FILE" 2>&1 || {
                warning "Failed to initialize some submodules, continuing anyway"
            }
            cd - > /dev/null
            
            log "Flash Attention cloned successfully!"
        else
            error "Flash Attention repository required. Please provide a valid path or allow cloning."
        fi
    else
        log "Flash Attention repository found at $FLASH_ATTN_PATH"
        # Update repository if it exists
        read -p "Would you like to update the Flash Attention repository? (y/n): " -n 1 -r
        echo
        if [[ $REPLY =~ ^[Yy]$ ]]; then
            log "Updating Flash Attention repository..."
            cd "$FLASH_ATTN_PATH"
            git pull >> "$LOG_FILE" 2>&1 || {
                warning "Failed to update repository, using existing version"
            }
            git submodule update --init --recursive >> "$LOG_FILE" 2>&1 || {
                warning "Failed to update submodules, continuing with existing"
            }
            cd - > /dev/null
        fi
    fi
    
    # Check if hopper directory exists (for FA v3)
    if [ ! -d "$FLASH_ATTN_PATH/hopper" ]; then
        warning "Flash Attention v3 hopper directory not found at $FLASH_ATTN_PATH/hopper"
        warning "This might be an older version of the repository"
        warning "Flash Attention v3 requires the hopper subdirectory for H100 optimization"
    fi
}

# Main setup function
main() {
    log "Starting Suno Environment Setup with Flash Attention v3 (H100 optimized)"
    log "Environment name: ${ENV_NAME}"
    log "Log file: ${LOG_FILE}"
    
    # Check CUDA version
    log "Checking CUDA version..."
    if command -v nvcc &> /dev/null; then
        cuda_version=$(nvcc --version | grep "release" | sed 's/.*release //' | cut -d',' -f1)
        log "CUDA version: ${cuda_version}"
        
        # Flash Attention v3 requires CUDA >= 12.3
        cuda_major=$(echo $cuda_version | cut -d'.' -f1)
        cuda_minor=$(echo $cuda_version | cut -d'.' -f2)
        
        if [ "$cuda_major" -lt 12 ] || ([ "$cuda_major" -eq 12 ] && [ "$cuda_minor" -lt 3 ]); then
            warning "Flash Attention v3 requires CUDA >= 12.3, found ${cuda_version}"
            warning "Installation may fail or performance may be suboptimal"
        fi
    else
        warning "nvcc not found. Unable to check CUDA version."
        warning "Flash Attention v3 requires CUDA >= 12.3"
    fi
    
    # Setup Suno Utils path
    setup_suno_utils_path
    
    # Setup Flash Attention path
    setup_flash_attention_path
    
    # Check if environment already exists
    if conda env list | grep -q "^${ENV_NAME} "; then
        warning "Environment ${ENV_NAME} already exists."
        read -p "Do you want to remove and recreate it? (y/n): " -n 1 -r
        echo
        if [[ $REPLY =~ ^[Yy]$ ]]; then
            log "Removing existing environment..."
            conda env remove -n "${ENV_NAME}" -y >> "$LOG_FILE" 2>&1
        else
            error "Environment already exists. Exiting."
        fi
    fi
    
    # Step 1: Create conda environment
    log "Step 1: Creating conda environment with Python 3.10..."
    conda create -n "${ENV_NAME}" python=3.10.15 -y >> "$LOG_FILE" 2>&1
    log "Conda environment created!"
    
    # Activate the environment
    source "$(conda info --base)/etc/profile.d/conda.sh"
    conda activate "${ENV_NAME}"
    
    # Step 2: Install suno_utils first
    log "Step 2: Installing suno_utils from ${SUNO_UTILS_PATH}..."
    pip install -e "$SUNO_UTILS_PATH" >> "$LOG_FILE" 2>&1
    log "suno_utils installed!"
    
    # Step 3: Check which torch version was installed
    log "Step 3: Checking torch version installed by suno_utils..."
    EXISTING_TORCH=$(python -c "import torch; print(torch.__version__)" 2>/dev/null || echo "None")
    log "Torch version installed by suno_utils: ${EXISTING_TORCH}"
    
    # Step 4: Uninstall torch (whatever version suno_utils installed)
    log "Step 4: Uninstalling torch installed by suno_utils..."
    pip uninstall torch torchvision torchaudio -y >> "$LOG_FILE" 2>&1
    log "Torch uninstalled!"
    
    # Step 5: Install PyTorch with CUDA 12.4 support
    log "Step 5: Installing PyTorch 2.6.0 with CUDA 12.4 support..."
    pip install torch==2.6.0 torchvision==0.21.0 torchaudio==2.6.0 --index-url https://download.pytorch.org/whl/cu124 >> "$LOG_FILE" 2>&1
    log "PyTorch installed with CUDA support!"
    
    # Verify PyTorch installation
    python -c "import torch; print(f'PyTorch version: {torch.__version__}')"
    python -c "import torch; print(f'CUDA available: {torch.cuda.is_available()}')"
    
    # Step 6: Build Flash Attention v3 from hopper directory
    log "Step 6: Building Flash Attention v3 (H100 optimized) from ${FLASH_ATTN_PATH}/hopper..."
    
    if [ ! -d "$FLASH_ATTN_PATH/hopper" ]; then
        warning "Flash Attention v3 hopper directory not found at $FLASH_ATTN_PATH/hopper"
        warning "Attempting to use Flash Attention v2 as fallback..."
        
        # Try regular Flash Attention v2
        cd "$FLASH_ATTN_PATH"
        
        # Clean any previous builds
        log "Cleaning previous Flash Attention builds..."
        rm -rf build dist *.egg-info >> "$LOG_FILE" 2>&1 || true
        
        # Install ninja
        log "Installing ninja..."
        pip install ninja >> "$LOG_FILE" 2>&1
        
        # Install Flash Attention v2
        log "Building and installing Flash Attention v2 as fallback (this may take 10-30 minutes)..."
        export MAX_JOBS=4  # Limit parallel jobs
        export FLASH_ATTENTION_FORCE_BUILD=TRUE
        pip install . --no-build-isolation >> "$LOG_FILE" 2>&1 || {
            warning "Flash Attention installation failed, continuing..."
        }
    else
        cd "$FLASH_ATTN_PATH/hopper"
        
        # Clean any previous builds
        log "Cleaning previous Flash Attention v3 builds..."
        rm -rf build dist *.egg-info >> "$LOG_FILE" 2>&1 || true
        
        # Install required packages
        log "Installing build dependencies..."
        pip install packaging ninja >> "$LOG_FILE" 2>&1
        
        # Install Flash Attention v3
        log "Building and installing Flash Attention v3 (this may take 10-30 minutes)..."
        log "Note: Flash Attention v3 is optimized for H100/H800 GPUs"
        
        export MAX_JOBS=4  # Limit parallel jobs
        export FLASH_ATTENTION_FORCE_BUILD=TRUE
        
        # Use python setup.py install as recommended in README
        python setup.py install >> "$LOG_FILE" 2>&1 || {
            warning "Flash Attention v3 installation failed"
            warning "This is expected if not running on H100/H800"
            warning "Trying Flash Attention v2 as fallback..."
            
            # Fallback to Flash Attention v2
            cd "$FLASH_ATTN_PATH"
            pip install . --no-build-isolation >> "$LOG_FILE" 2>&1 || {
                warning "Flash Attention v2 installation also failed, continuing without Flash Attention..."
            }
        }
    fi
    
    # Test Flash Attention import
    log "Testing Flash Attention installation..."
    python -c "
try:
    import flash_attn_interface
    print('✓ Flash Attention v3 successfully installed')
    print(f'  Module location: {flash_attn_interface.__file__}')
except ImportError:
    try:
        import flash_attn
        print('✓ Flash Attention v2 installed as fallback')
    except ImportError:
        print('✗ Flash Attention not available')
" | tee -a "$LOG_FILE"
    
    # Step 7: Install additional Python packages
    log "Step 7: Installing additional Python packages..."
    pip install wandb nnAudio deepspeed auraloss torchsde g2p_en \
        transformers==4.44.0 modal==1.0.1 better_profanity encodec \
        pytorch-lightning sentencepiece tiktoken einops ffmpeg-python >> "$LOG_FILE" 2>&1
    log "Additional packages installed!"
    
    # Step 8: Install sox via conda
    log "Step 8: Installing sox via conda..."
    conda install -c conda-forge sox -y >> "$LOG_FILE" 2>&1
    log "Sox installed!"
    
    # Return to script directory
    cd "${SCRIPT_DIR}"
    
    # Verification
    log "Running verification tests..."
    python "${SCRIPT_DIR}/verify_imports.py" | tee -a "$LOG_FILE"
    
    # Create activation script
    log "Creating activation script..."
    cat > "${SCRIPT_DIR}/activate_${ENV_NAME}.sh" << EOF
#!/bin/bash
# Activation script for ${ENV_NAME} environment with Flash Attention v3

# Source conda
source "\$(conda info --base)/etc/profile.d/conda.sh"

# Activate environment
conda activate ${ENV_NAME}

# Set environment variables
export CUDA_VISIBLE_DEVICES=\${CUDA_VISIBLE_DEVICES:-0}
export OMP_NUM_THREADS=1
export TRITON_CACHE_DIR=/mnt/localdisk/.triton_cache_\$USER
export PYTHONPATH="${FLASH_ATTN_PATH}/hopper:\$PYTHONPATH"

echo "Environment ${ENV_NAME} activated (Flash Attention v3)!"
echo "Python: \$(which python)"
echo "PyTorch: \$(python -c 'import torch; print(torch.__version__)' 2>/dev/null || echo 'Not available')"
echo "CUDA available: \$(python -c 'import torch; print(torch.cuda.is_available())' 2>/dev/null || echo 'Unknown')"
echo "Flash Attention path: ${FLASH_ATTN_PATH}"
echo "Suno Utils path: ${SUNO_UTILS_PATH}"

# Test Flash Attention version
python -c "
try:
    import flash_attn_interface
    print('Flash Attention: v3 (H100 optimized)')
except ImportError:
    try:
        import flash_attn
        print('Flash Attention: v2 (fallback)')
    except ImportError:
        print('Flash Attention: Not available')
" 2>/dev/null
EOF
    
    chmod +x "${SCRIPT_DIR}/activate_${ENV_NAME}.sh"
    
    log "================================================"
    log "Environment setup completed!"
    log "Flash Attention path: ${FLASH_ATTN_PATH}"
    log "Suno Utils path: ${SUNO_UTILS_PATH}"
    log ""
    log "To activate the environment, run:"
    log "  source ${SCRIPT_DIR}/activate_${ENV_NAME}.sh"
    log "Or:"
    log "  conda activate ${ENV_NAME}"
    log ""
    log "Note: Flash Attention v3 is optimized for H100/H800 GPUs"
    log "================================================"
}

# Run main function
main "$@"