#!/bin/bash

# Download Essentia models for audio classification
# Models from: https://essentia.upf.edu/models.html

set -e  # Exit on error

MODELS_DIR="backend/models"
BASE_URL="https://essentia.upf.edu/models/classification-heads"

echo "📦 Downloading Essentia models..."
echo "Models directory: $MODELS_DIR"

# Create models directory if it doesn't exist
mkdir -p "$MODELS_DIR"

# Model files
declare -A MODELS
MODELS=(
    ["mtg_jamendo_genre-discogs-effnet-1.pb"]="$BASE_URL/mtg_jamendo_genre/mtg_jamendo_genre-discogs-effnet-1.pb"
    ["mtg_jamendo_moodtheme-discogs-effnet-1.pb"]="$BASE_URL/mtg_jamendo_moodtheme/mtg_jamendo_moodtheme-discogs-effnet-1.pb"
    ["mtg_jamendo_instrument-discogs-effnet-1.pb"]="$BASE_URL/mtg_jamendo_instrument/mtg_jamendo_instrument-discogs-effnet-1.pb"
)

# Download each model
for model_file in "${!MODELS[@]}"; do
    url="${MODELS[$model_file]}"
    output_path="$MODELS_DIR/$model_file"

    if [ -f "$output_path" ]; then
        echo "✓ $model_file already exists, skipping..."
    else
        echo "⬇️  Downloading $model_file..."
        curl -L -o "$output_path" "$url"

        if [ -f "$output_path" ]; then
            echo "✓ Downloaded $model_file"
        else
            echo "✗ Failed to download $model_file"
            exit 1
        fi
    fi
done

echo ""
echo "✅ All models downloaded successfully!"
echo ""
echo "Models available:"
ls -lh "$MODELS_DIR"/*.pb 2>/dev/null || echo "No .pb files found"

echo ""
echo "Note: Class labels are defined in backend/src/core/essentia_classifier.py"
echo "You can now start the backend with: docker-compose up"