#!/bin/bash

# Robust Virtualenv Setup Script for Suno GPT
# Alternative to conda environment using virtualenv/venv
# Supports both standard pip and uv for faster installation

set -e  # Exit on error

# Configuration
ENV_NAME="${1:-suno_venv}"
PYTHON_VERSION="python3.10"
CUDA_VERSION="cu124"
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
ENV_DIR="${SCRIPT_DIR}/${ENV_NAME}"
LOG_FILE="${SCRIPT_DIR}/setup_venv_${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'
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"
}

# Check if UV is available for faster package installation
check_uv() {
    if command -v uv &> /dev/null; then
        log "UV detected - will use for faster package installation"
        USE_UV=true
    else
        info "UV not found - using standard pip (consider installing UV for 10-100x faster installation)"
        info "Install UV with: curl -LsSf https://astral.sh/uv/install.sh | sh"
        USE_UV=false
    fi
}

# Check prerequisites
check_prerequisites() {
    log "Checking prerequisites..."
    
    # Check Python version
    if ! command -v ${PYTHON_VERSION} &> /dev/null; then
        # Try python3 as fallback
        if command -v python3 &> /dev/null; then
            PYTHON_CMD="python3"
            PYTHON_VER=$(python3 --version | cut -d' ' -f2)
            if [[ ! "$PYTHON_VER" =~ ^3\.10 ]]; then
                error "Python 3.10 required, found ${PYTHON_VER}"
            fi
        else
            error "${PYTHON_VERSION} not found. Please install Python 3.10"
        fi
    else
        PYTHON_CMD=${PYTHON_VERSION}
    fi
    
    log "Using Python: $(${PYTHON_CMD} --version)"
    
    # 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 [ -d "$ENV_DIR" ]; then
        warning "Environment directory '${ENV_DIR}' 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..."
            rm -rf "$ENV_DIR"
        else
            error "Environment already exists. Exiting."
        fi
    fi
    
    # Check disk space
    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
    
    # Check for required system packages
    MISSING_PACKAGES=""
    for pkg in gcc g++ make cmake; do
        if ! command -v $pkg &> /dev/null; then
            MISSING_PACKAGES="$MISSING_PACKAGES $pkg"
        fi
    done
    
    if [ -n "$MISSING_PACKAGES" ]; then
        warning "Missing system packages:$MISSING_PACKAGES"
        info "These may be needed for building some packages"
    fi
    
    check_uv
    log "Prerequisites check passed!"
}

# Create virtual environment
create_virtualenv() {
    log "Creating virtual environment at ${ENV_DIR}..."
    
    ${PYTHON_CMD} -m venv "$ENV_DIR" || error "Failed to create virtual environment"
    
    # Activate environment
    source "${ENV_DIR}/bin/activate"
    
    # Upgrade pip, setuptools, wheel
    log "Upgrading pip, setuptools, and wheel..."
    pip install --upgrade pip setuptools wheel >> "$LOG_FILE" 2>&1 || error "Failed to upgrade pip"
    
    log "Virtual environment created successfully!"
}

# Install packages using UV
install_with_uv() {
    log "Installing packages with UV (fast mode)..."
    
    source "${ENV_DIR}/bin/activate"
    
    # Install UV in the environment if not globally available
    pip install uv >> "$LOG_FILE" 2>&1
    
    # 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}..."
    uv 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 with UV, trying latest..."
        uv pip install torch torchvision torchaudio \
            --index-url https://download.pytorch.org/whl/${CUDA_VERSION} >> "$LOG_FILE" 2>&1 || \
            error "Failed to install PyTorch"
    }
    
    # Install other packages from requirements (skip torch since we already installed it)
    log "Installing remaining packages..."
    grep -v "^torch==" "${SCRIPT_DIR}/requirements.txt" | grep -v "^torchvision==" | grep -v "^torchaudio==" > /tmp/requirements_no_torch.txt
    uv pip install -r /tmp/requirements_no_torch.txt >> "$LOG_FILE" 2>&1 || \
        warning "Some packages failed to install"
    rm /tmp/requirements_no_torch.txt
}

# Install packages using pip
install_with_pip() {
    log "Installing packages with pip (standard mode)..."
    
    source "${ENV_DIR}/bin/activate"
    
    # 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"
    }
    
    # Install packages in batches to avoid dependency conflicts
    log "Installing ML frameworks..."
    pip install \
        transformers==4.44.0 \
        wandb==0.21.0 \
        deepspeed==0.17.2 \
        pytorch-lightning==2.5.2 \
        >> "$LOG_FILE" 2>&1 || warning "Some ML packages failed"
    
    log "Installing audio 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"
    
    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"
    
    log "Installing utilities..."
    pip install \
        einops==0.8.1 \
        torchsde==0.2.6 \
        modal==1.1.3 \
        ninja==1.11.1.4 \
        >> "$LOG_FILE" 2>&1 || warning "Some utility packages failed"
}

# Install Flash Attention
install_flash_attention() {
    log "Installing Flash Attention v2..."
    
    source "${ENV_DIR}/bin/activate"
    
    # Try pip installation first (pre-built wheels if available)
    log "Attempting to install Flash Attention from pip..."
    pip install flash-attn==2.8.1 --no-build-isolation >> "$LOG_FILE" 2>&1 && {
        log "Flash Attention installed from pip successfully!"
        return 0
    }
    
    warning "Pip installation failed, building from source..."
    
    # Clone and build from source
    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 || {
            warning "Failed to clone Flash Attention, skipping..."
            return 1
        }
    fi
    
    cd "$FLASH_DIR"
    
    # Set build environment
    export MAX_JOBS=4
    export FLASH_ATTENTION_FORCE_BUILD=TRUE
    
    log "Building Flash Attention (this may take 10-30 minutes)..."
    pip install . >> "$LOG_FILE" 2>&1 || {
        warning "Failed to build Flash Attention from source"
        return 1
    }
    
    cd "$SCRIPT_DIR"
    log "Flash Attention built from source successfully!"
}

# Install suno_utils
install_suno_utils() {
    log "Installing suno_utils..."
    
    source "${ENV_DIR}/bin/activate"
    
    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
    
    pip install -e "$SUNO_UTILS_PATH" >> "$LOG_FILE" 2>&1 || error "Failed to install suno_utils"
    
    log "suno_utils installed successfully!"
}

# Install system dependencies
install_system_deps() {
    log "Checking for sox installation..."
    
    if ! command -v sox &> /dev/null; then
        warning "sox not found in system"
        info "To install sox:"
        info "  Ubuntu/Debian: sudo apt-get install sox libsox-fmt-all"
        info "  CentOS/RHEL: sudo yum install sox"
        info "  macOS: brew install sox"
    else
        log "sox found: $(which sox)"
    fi
}

# Verify installation
verify_installation() {
    log "Verifying installation..."
    
    source "${ENV_DIR}/bin/activate"
    
    # Create verification script
    cat > "${SCRIPT_DIR}/verify_venv_imports.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"),
    ("einops", "Einops"),
    ("torchsde", "Torch SDE"),
]

print("Testing package imports...")
print("-" * 40)

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

print("-" * 40)

# Test CUDA
try:
    import torch
    cuda_available = torch.cuda.is_available()
    cuda_version = torch.version.cuda if cuda_available else "N/A"
    print(f"CUDA Available: {cuda_available}")
    print(f"CUDA Version: {cuda_version}")
    print(f"PyTorch Version: {torch.__version__}")
except:
    print("Could not check CUDA status")

if failed:
    print(f"\n⚠ Warning: {len(failed)} packages failed to import")
    sys.exit(1)
else:
    print("\n✓ All packages imported successfully!")
EOF
    
    python "${SCRIPT_DIR}/verify_venv_imports.py" || warning "Some packages failed verification"
    
    log "Installation verification complete!"
}

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

# Activate virtual environment
source "${ENV_DIR}/bin/activate"

# 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 "Virtual environment ${ENV_NAME} activated!"
echo "Python: \$(which python)"
echo "Python version: \$(python --version)"
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 ""
echo "To deactivate, run: deactivate"
EOF
    
    chmod +x "${SCRIPT_DIR}/activate_${ENV_NAME}.sh"
    
    # Quick test script
    cat > "${SCRIPT_DIR}/test_${ENV_NAME}.py" << 'EOF'
#!/usr/bin/env python
"""Quick test script for the virtual environment"""

import sys
print(f"Python: {sys.version}")
print(f"Executable: {sys.executable}")

try:
    import torch
    print(f"PyTorch: {torch.__version__}")
    print(f"CUDA available: {torch.cuda.is_available()}")
    if torch.cuda.is_available():
        print(f"CUDA device: {torch.cuda.get_device_name(0)}")
except ImportError:
    print("PyTorch not available")

try:
    import transformers
    print(f"Transformers: {transformers.__version__}")
except ImportError:
    print("Transformers not available")

try:
    import flash_attn
    print(f"Flash Attention: {flash_attn.__version__}")
except ImportError:
    print("Flash Attention not available")

try:
    import suno_utils
    print("Suno Utils: available")
except ImportError:
    print("Suno Utils not available")
EOF
    
    chmod +x "${SCRIPT_DIR}/test_${ENV_NAME}.py"
    
    log "Scripts created:"
    log "  - Activation: ${SCRIPT_DIR}/activate_${ENV_NAME}.sh"
    log "  - Quick test: ${SCRIPT_DIR}/test_${ENV_NAME}.py"
}

# Main execution
main() {
    log "Starting Virtual Environment setup for ${ENV_NAME}"
    log "Log file: ${LOG_FILE}"
    
    check_prerequisites
    create_virtualenv
    
    # Install suno_utils first (it will pull in its torch dependency)
    install_suno_utils
    
    # Then install/upgrade packages including CUDA-enabled PyTorch
    if [ "$USE_UV" = true ]; then
        install_with_uv
    else
        install_with_pip
    fi
    
    install_flash_attention
    install_system_deps
    verify_installation
    create_activation_script
    
    log "================================================"
    log "Virtual environment setup completed!"
    log ""
    log "To activate the environment, run:"
    log "  source ${SCRIPT_DIR}/activate_${ENV_NAME}.sh"
    log "Or:"
    log "  source ${ENV_DIR}/bin/activate"
    log ""
    log "To test the environment:"
    log "  python ${SCRIPT_DIR}/test_${ENV_NAME}.py"
    log "================================================"
}

# Run main function
main "$@"