#!/bin/bash

# Robust Conda Environment Setup Script for Suno GPT
# This script creates a conda environment with all necessary dependencies
# including PyTorch, Flash Attention v2, and audio processing tools

set -e  # Exit on error

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

# Colors for output
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
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"
}

# Check prerequisites
check_prerequisites() {
    log "Checking prerequisites..."
    
    # Check conda is available
    if ! command -v conda &> /dev/null; then
        error "Conda is not installed or not in PATH"
    fi
    
    # Check CUDA availability
    if ! nvidia-smi &> /dev/null; then
        warning "NVIDIA GPU not detected. CUDA packages will still be installed but may not work properly."
    else
        log "NVIDIA GPU detected: $(nvidia-smi --query-gpu=name --format=csv,noheader | head -1)"
    fi
    
    # 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
        else
            error "Environment already exists. Exiting."
        fi
    fi
    
    # Check disk space (need at least 20GB)
    AVAILABLE_SPACE=$(df "${SCRIPT_DIR}" | awk 'NR==2 {print int($4/1024/1024)}')
    if [ "$AVAILABLE_SPACE" -lt 20 ]; then
        error "Insufficient disk space. At least 20GB required, only ${AVAILABLE_SPACE}GB available."
    fi
    
    log "Prerequisites check passed!"
}

# Create conda environment
create_conda_env() {
    log "Creating conda environment '${ENV_NAME}' with Python ${PYTHON_VERSION}..."
    
    # Try to create with specific Python version, fallback to 3.10 if exact version not available
    if ! conda create -n "${ENV_NAME}" python="${PYTHON_VERSION}" -y >> "$LOG_FILE" 2>&1; then
        warning "Python ${PYTHON_VERSION} not available, trying Python 3.10..."
        conda create -n "${ENV_NAME}" python=3.10 -y >> "$LOG_FILE" 2>&1 || error "Failed to create conda environment"
    fi
    
    log "Conda environment created successfully!"
}

# Install PyTorch with CUDA support
install_pytorch() {
    log "Installing PyTorch with CUDA ${CUDA_VERSION} support..."
    
    # Activate environment for this subshell
    source "$(conda info --base)/etc/profile.d/conda.sh"
    conda activate "${ENV_NAME}"
    
    # Check what torch version suno_utils installed
    if python -c "import torch" 2>/dev/null; then
        EXISTING_VERSION=$(python -c "import torch; print(torch.__version__)")
        log "suno_utils installed PyTorch ${EXISTING_VERSION}"
        log "Uninstalling to replace with CUDA-enabled version..."
    fi
    
    # Always uninstall existing torch to ensure clean CUDA installation
    pip uninstall torch torchvision torchaudio -y >> "$LOG_FILE" 2>&1 || true
    
    # Install PyTorch with CUDA support
    log "Installing PyTorch 2.6.0 with CUDA ${CUDA_VERSION}..."
    pip install torch==2.6.0 torchvision==0.21.0 torchaudio==2.6.0 \
        --index-url https://download.pytorch.org/whl/${CUDA_VERSION} >> "$LOG_FILE" 2>&1 || {
        warning "Failed to install specific PyTorch version, trying latest..."
        pip install torch torchvision torchaudio \
            --index-url https://download.pytorch.org/whl/${CUDA_VERSION} >> "$LOG_FILE" 2>&1 || \
            error "Failed to install PyTorch"
    }
    
    # Verify PyTorch installation
    python -c "import torch; print(f'PyTorch {torch.__version__} installed')" || error "PyTorch installation verification failed"
    python -c "import torch; print(f'CUDA available: {torch.cuda.is_available()}')" || warning "CUDA not available in PyTorch"
    
    log "PyTorch installed successfully!"
}

# Install Flash Attention v2
install_flash_attention() {
    log "Installing Flash Attention v2..."
    
    source "$(conda info --base)/etc/profile.d/conda.sh"
    conda activate "${ENV_NAME}"
    
    # Install ninja for compilation
    pip install ninja >> "$LOG_FILE" 2>&1 || error "Failed to install ninja"
    
    # Clone Flash Attention if not already present
    FLASH_DIR="${SCRIPT_DIR}/flash-attention"
    if [ ! -d "$FLASH_DIR" ]; then
        log "Cloning Flash Attention repository..."
        git clone https://github.com/Dao-AILab/flash-attention.git "$FLASH_DIR" >> "$LOG_FILE" 2>&1 || \
            error "Failed to clone Flash Attention repository"
    else
        log "Flash Attention repository already exists, pulling latest..."
        cd "$FLASH_DIR"
        git pull >> "$LOG_FILE" 2>&1 || warning "Failed to update Flash Attention repository"
    fi
    
    # Build and install Flash Attention
    cd "$FLASH_DIR"
    log "Building Flash Attention (this may take 10-30 minutes)..."
    
    # Set environment variables for compilation
    export MAX_JOBS=4  # Limit parallel jobs to avoid OOM
    export FLASH_ATTENTION_FORCE_BUILD=TRUE
    
    python setup.py install >> "$LOG_FILE" 2>&1 || {
        warning "Failed to build Flash Attention from source, trying pip..."
        pip install flash-attn==2.8.1 --no-build-isolation >> "$LOG_FILE" 2>&1 || \
            warning "Flash Attention installation failed, continuing without it"
    }
    
    # Verify Flash Attention installation
    python -c "import flash_attn; print(f'Flash Attention {flash_attn.__version__} installed')" 2>/dev/null || \
        warning "Flash Attention not available"
    
    cd "$SCRIPT_DIR"
}

# Install suno_utils
install_suno_utils() {
    log "Installing suno_utils..."
    
    source "$(conda info --base)/etc/profile.d/conda.sh"
    conda activate "${ENV_NAME}"
    
    SUNO_UTILS_PATH="/home/vibert/projects/glockenspiel/suno_utils"
    
    if [ ! -d "$SUNO_UTILS_PATH" ]; then
        error "suno_utils not found at ${SUNO_UTILS_PATH}"
    fi
    
    cd "$SUNO_UTILS_PATH"
    pip install -e . >> "$LOG_FILE" 2>&1 || error "Failed to install suno_utils"
    
    # Verify installation
    python -c "import suno_utils; print('suno_utils installed')" || error "suno_utils installation verification failed"
    
    cd "$SCRIPT_DIR"
    log "suno_utils installed successfully!"
}

# Install additional packages
install_additional_packages() {
    log "Installing additional packages..."
    
    source "$(conda info --base)/etc/profile.d/conda.sh"
    conda activate "${ENV_NAME}"
    
    # Install core ML packages
    log "Installing ML frameworks..."
    pip install \
        wandb==0.21.0 \
        transformers==4.44.0 \
        deepspeed==0.17.2 \
        pytorch-lightning==2.5.2 \
        >> "$LOG_FILE" 2>&1 || warning "Some ML packages failed to install"
    
    # Install audio processing packages
    log "Installing audio processing packages..."
    pip install \
        nnAudio==0.3.3 \
        auraloss==0.4.0 \
        encodec==0.1.1 \
        descript-audiotools==0.7.1 \
        >> "$LOG_FILE" 2>&1 || warning "Some audio packages failed to install"
    
    # Install NLP packages
    log "Installing NLP packages..."
    pip install \
        g2p-en==2.1.0 \
        phonemizer==3.2.1 \
        sentencepiece==0.1.97 \
        tiktoken==0.1.2 \
        better-profanity==0.7.0 \
        >> "$LOG_FILE" 2>&1 || warning "Some NLP packages failed to install"
    
    # Install utilities
    log "Installing utility packages..."
    pip install \
        einops==0.8.1 \
        torchsde==0.2.6 \
        modal==1.1.3 \
        >> "$LOG_FILE" 2>&1 || warning "Some utility packages failed to install"
    
    # Install sox via conda
    log "Installing sox via conda..."
    conda install -c conda-forge sox -y >> "$LOG_FILE" 2>&1 || warning "Failed to install sox"
    
    log "Additional packages installed!"
}

# Verify installation
verify_installation() {
    log "Verifying installation..."
    
    source "$(conda info --base)/etc/profile.d/conda.sh"
    conda activate "${ENV_NAME}"
    
    # Create verification script
    cat > "${SCRIPT_DIR}/verify_imports_temp.py" << 'EOF'
import sys
import importlib

packages_to_test = [
    ("torch", "PyTorch"),
    ("transformers", "Transformers"),
    ("wandb", "Weights & Biases"),
    ("deepspeed", "DeepSpeed"),
    ("nnAudio", "nnAudio"),
    ("auraloss", "Auraloss"),
    ("flash_attn", "Flash Attention"),
    ("suno_utils", "Suno Utils"),
    ("g2p_en", "G2P English"),
    ("modal", "Modal"),
]

failed = []
for package, name in packages_to_test:
    try:
        importlib.import_module(package)
        print(f"✓ {name}")
    except ImportError as e:
        print(f"✗ {name}: {e}")
        failed.append(name)

if failed:
    print(f"\nWarning: Failed to import: {', '.join(failed)}")
    sys.exit(1)
else:
    print("\nAll packages imported successfully!")
EOF
    
    python "${SCRIPT_DIR}/verify_imports_temp.py" || warning "Some packages failed to import"
    rm "${SCRIPT_DIR}/verify_imports_temp.py"
    
    log "Installation verification complete!"
}

# Create activation script
create_activation_script() {
    log "Creating activation script..."
    
    cat > "${SCRIPT_DIR}/activate_${ENV_NAME}.sh" << EOF
#!/bin/bash
# Activation script for ${ENV_NAME} environment

# 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

echo "Environment ${ENV_NAME} activated!"
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')"
EOF
    
    chmod +x "${SCRIPT_DIR}/activate_${ENV_NAME}.sh"
    log "Activation script created: ${SCRIPT_DIR}/activate_${ENV_NAME}.sh"
}

# Main execution
main() {
    log "Starting Conda environment setup for ${ENV_NAME}"
    log "Log file: ${LOG_FILE}"
    
    check_prerequisites
    create_conda_env
    install_suno_utils
    install_pytorch
    install_flash_attention
    install_additional_packages
    verify_installation
    create_activation_script
    
    log "================================================"
    log "Environment setup completed successfully!"
    log "To activate the environment, run:"
    log "  source ${SCRIPT_DIR}/activate_${ENV_NAME}.sh"
    log "Or:"
    log "  conda activate ${ENV_NAME}"
    log "================================================"
}

# Run main function
main "$@"