#!/bin/bash

# setup_musicgen_env.sh
# Automates the setup of a clean virtual environment for MusicGen script
# Handles dependency conflicts, installs compatible versions, and verifies setup
# Designed for Ubuntu homelab (T-1000) with RTX 3060 Ti, CUDA 12, Python 3.10

set -e  # Exit on any error

# Define variables
VENV_DIR="$HOME/ghostai_music_generator/musicgen_env"
LOG_FILE="$HOME/ghostai_music_generator/setup_musicgen_env_$(date +%Y%m%d_%H%M%S).log"
PYTHON_VERSION="python3.10"
PIP="$VENV_DIR/bin/pip"
PYTHON="$VENV_DIR/bin/python"

# Create log directory
mkdir -p "$(dirname "$LOG_FILE")"

# Function to log messages
log() {
    echo "[$(date +%Y-%m-%d\ %H:%M:%S)] $1" | tee -a "$LOG_FILE"
}

# Function to check command status
check_status() {
    if [ $? -ne 0 ]; then
        log "ERROR: $1 failed"
        exit 1
    fi
}

# Function to check if a package is installed
check_package() {
    dpkg -l "$1" &>/dev/null
    return $?
}

# Start logging
log "Starting MusicGen environment setup"

# Install system dependencies
log "Installing system dependencies..."
sudo apt-get update >> "$LOG_FILE" 2>&1
check_status "apt-get update"

for pkg in build-essential cmake python3.10 python3.10-venv python3.10-dev; do
    if ! check_package "$pkg"; then
        log "Installing $pkg..."
        sudo apt-get install -y "$pkg" >> "$LOG_FILE" 2>&1
        check_status "install $pkg"
    else
        log "$pkg already installed"
    fi
done

# Remove existing virtual environment if it exists
if [ -d "$VENV_DIR" ]; then
    log "Removing existing virtual environment at $VENV_DIR"
    rm -rf "$VENV_DIR"
    check_status "remove virtual environment"
fi

# Create and activate virtual environment
log "Creating virtual environment at $VENV_DIR"
$PYTHON_VERSION -m venv "$VENV_DIR"
check_status "create virtual environment"
source "$VENV_DIR/bin/activate"
check_status "activate virtual environment"

# Upgrade pip
log "Upgrading pip..."
$PIP install --upgrade pip >> "$LOG_FILE" 2>&1
check_status "upgrade pip"

# Uninstall conflicting packages
log "Uninstalling conflicting packages..."
$PIP uninstall -y torch torchaudio numpy transformers requests spacy networkx audiocraft pydub gradio typer pydantic laion-clap pydantic-core xformers torchdata torchtext torchvision stable-audio-tools accelerate dctorch >> "$LOG_FILE" 2>&1
check_status "uninstall conflicting packages"

# Install dependencies
log "Installing dependencies..."
$PIP install torch==2.1.0+cu121 torchaudio==2.1.0+cu121 --index-url https://download.pytorch.org/whl/cu121 >> "$LOG_FILE" 2>&1
check_status "install torch and torchaudio"
$PIP install git+https://github.com/facebookresearch/audiocraft.git@refs/tags/v1.3.0 numpy==1.26.4 transformers==4.40.2 requests==2.31.0 spacy==3.7.2 networkx==2.8.8 pydub==0.25.1 gradio==3.50.2 pydantic==1.10.13 >> "$LOG_FILE" 2>&1
check_status "install remaining dependencies"

# Verify installation
log "Verifying installed packages..."
$PIP list | grep -E "torch|torchaudio|numpy|transformers|requests|spacy|networkx|audiocraft|pydub|gradio|typer|pydantic|laion-clap" > "$HOME/ghostai_music_generator/verified_packages.txt"
check_status "list installed packages"

# Check CUDA and torch
log "Verifying CUDA and torch..."
$PYTHON -c "import torch; print(torch.__version__, torch.cuda.is_available(), torch.version.cuda)" >> "$LOG_FILE" 2>&1
check_status "verify CUDA and torch"

# Test imports
log "Testing imports..."
$PYTHON -c "import torch; import torchaudio; import audiocraft; import numpy; import transformers; import requests; import spacy; import networkx; import pydub; import gradio; import pydantic; print('All imports successful')" >> "$LOG_FILE" 2>&1
check_status "test imports"

# Check for dependency conflicts
log "Checking for dependency conflicts with pipdeptree..."
$PIP install pipdeptree >> "$LOG_FILE" 2>&1
check_status "install pipdeptree"
$PYTHON -m pipdeptree > "$HOME/ghostai_music_generator/pipdeptree_output.txt" 2>> "$LOG_FILE"
check_status "run pipdeptree"

# Clean up unnecessary packages
log "Cleaning up unnecessary packages..."
$PIP uninstall -y stable-audio-tools laion-clap x-transformers accelerate dctorch torchdata torchtext torchvision >> "$LOG_FILE" 2>&1 || true  # Ignore errors if not installed

# Final verification
log "Final verification of installed packages..."
EXPECTED_PACKAGES=(
    "audiocraft==1.3.0"
    "gradio==3.50.2"
    "networkx==2.8.8"
    "numpy==1.26.4"
    "pydantic==1.10.13"
    "pydub==0.25.1"
    "requests==2.31.0"
    "spacy==3.7.2"
    "torch==2.1.0+cu121"
    "torchaudio==2.1.0+cu121"
    "transformers==4.40.2"
    "typer==0.7.0"
)
for pkg in "${EXPECTED_PACKAGES[@]}"; do
    if $PIP list | grep -q "$pkg"; then
        log "$pkg verified"
    else
        log "ERROR: $pkg not found"
        exit 1
    fi
done

log "MusicGen environment setup completed successfully"
log "Virtual environment: $VENV_DIR"
log "Log file: $LOG_FILE"
log "Verified packages: $HOME/ghostai_music_generator/verified_packages.txt"
log "Dependency tree: $HOME/ghostai_music_generator/pipdeptree_output.txt"
log "To activate: source $VENV_DIR/bin/activate"
log "To run script: python musicgen_optimized.py"

exit 0