#!/bin/bash
set -e

echo "==================================="
echo "Audio Separator API - Install Script"
echo "==================================="

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

print_status() {
    echo -e "${GREEN}[OK]${NC} $1"
}

print_warning() {
    echo -e "${YELLOW}[WARN]${NC} $1"
}

print_error() {
    echo -e "${RED}[ERROR]${NC} $1"
}

# Check if running as root
if [ "$EUID" -eq 0 ]; then
    print_warning "Running as root - this is fine for VM setup"
fi

# Detect OS
OS="unknown"
if [ -f /etc/os-release ]; then
    . /etc/os-release
    OS=$ID
fi

echo ""
echo "Detected OS: $OS"
echo ""

# Step 1: Install system dependencies
echo "Step 1: Installing system dependencies..."

if [ "$OS" = "ubuntu" ] || [ "$OS" = "debian" ]; then
    apt-get update
    apt-get install -y \
        python3.10 \
        python3.10-venv \
        python3-pip \
        ffmpeg \
        libsndfile1 \
        curl \
        git
    print_status "System dependencies installed"

elif [ "$OS" = "centos" ] || [ "$OS" = "rhel" ] || [ "$OS" = "fedora" ]; then
    dnf install -y \
        python3.10 \
        python3-pip \
        ffmpeg \
        libsndfile \
        curl \
        git
    print_status "System dependencies installed"

elif [[ "$OSTYPE" == "darwin"* ]]; then
    if ! command -v brew &> /dev/null; then
        print_error "Homebrew not found. Please install it first."
        exit 1
    fi
    brew install python@3.10 ffmpeg libsndfile
    print_status "System dependencies installed"

else
    print_warning "Unknown OS - please ensure python3.10, ffmpeg, and libsndfile are installed"
fi

# Step 2: Install uv (fast Python package manager)
echo ""
echo "Step 2: Installing uv package manager..."

if ! command -v uv &> /dev/null; then
    curl -LsSf https://astral.sh/uv/install.sh | sh
    export PATH="$HOME/.local/bin:$PATH"
    print_status "uv installed"
else
    print_status "uv already installed"
fi

# Step 3: Create virtual environment and install dependencies
echo ""
echo "Step 3: Setting up Python environment..."

cd "$(dirname "$0")"

# Remove existing venv if present
if [ -d ".venv" ]; then
    rm -rf .venv
fi

uv venv --python 3.10
print_status "Virtual environment created"

# Step 4: Install Python dependencies
echo ""
echo "Step 4: Installing Python dependencies..."

# Install with CUDA support detection
uv pip install -e ".[api]"
print_status "Base dependencies installed"

# Step 5: Check for NVIDIA GPU and install CUDA dependencies
echo ""
echo "Step 5: Checking GPU availability..."

if command -v nvidia-smi &> /dev/null; then
    echo ""
    nvidia-smi
    echo ""
    print_status "NVIDIA GPU detected"

    # Install PyTorch with CUDA
    echo "Installing PyTorch with CUDA support..."
    uv pip install torch torchvision --index-url https://download.pytorch.org/whl/cu121
    print_status "PyTorch with CUDA installed"
else
    print_warning "No NVIDIA GPU detected - will use CPU (slower)"
fi

# Step 6: Pre-download default model
echo ""
echo "Step 6: Pre-downloading default model (this may take a few minutes)..."

.venv/bin/python -c "
from audio_separator.separator import Separator
import os
os.makedirs('/tmp/audio-separator/models', exist_ok=True)
s = Separator(model_file_dir='/tmp/audio-separator/models')
s.load_model()
print('Model downloaded successfully')
"
print_status "Default model downloaded"

# Step 7: Create systemd service file (optional)
echo ""
echo "Step 7: Creating systemd service file..."

SERVICE_FILE="/etc/systemd/system/audio-separator.service"
INSTALL_DIR="$(pwd)"

if [ -d "/etc/systemd/system" ] && [ "$EUID" -eq 0 ]; then
    cat > "$SERVICE_FILE" << EOF
[Unit]
Description=Audio Separator API
After=network.target

[Service]
Type=simple
User=root
WorkingDirectory=$INSTALL_DIR
ExecStart=$INSTALL_DIR/.venv/bin/uvicorn app:app --host 0.0.0.0 --port 8000
Restart=always
RestartSec=10
Environment="PATH=$INSTALL_DIR/.venv/bin:/usr/local/bin:/usr/bin:/bin"

[Install]
WantedBy=multi-user.target
EOF
    systemctl daemon-reload
    print_status "Systemd service created at $SERVICE_FILE"
    echo "    To enable: systemctl enable audio-separator"
    echo "    To start:  systemctl start audio-separator"
else
    print_warning "Skipping systemd service (not root or systemd not available)"
fi

echo ""
echo "==================================="
echo "Installation complete!"
echo "==================================="
echo ""
echo "To run the API:"
echo "  .venv/bin/uvicorn app:app --host 0.0.0.0 --port 8000"
echo ""
echo "To run tests:"
echo "  ./test.sh"
echo ""
echo "API will be available at http://localhost:8000"
echo "API docs at http://localhost:8000/docs"
echo ""