#!/usr/bin/env bash

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

set -e  # Exit on error

MODELS_DIR="backend/models"
CLASS_HEADS_URL="https://essentia.upf.edu/models/classification-heads"
EMBEDDINGS_URL="https://essentia.upf.edu/models/feature-extractors/discogs-effnet"

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

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

# Download function
download_model() {
    local model_file="$1"
    local url="$2"
    local output_path="$MODELS_DIR/$model_file"

    if [ -f "$output_path" ]; then
        echo "✓ $model_file already exists, skipping..."
    else
        echo "⬇️  Downloading $model_file..."
        # Use -k flag to ignore SSL certificate issues with essentia.upf.edu
        curl -k -L -o "$output_path" "$url"

        if [ -f "$output_path" ] && [ -s "$output_path" ]; then
            echo "✓ Downloaded $model_file ($(du -h "$output_path" | cut -f1))"
        else
            echo "✗ Failed to download $model_file"
            rm -f "$output_path"  # Remove empty/failed file
            exit 1
        fi
    fi
}

# Download embedding model first (required for all classification heads)
echo ""
echo "Downloading embedding model..."
download_model "discogs-effnet-bs64-1.pb" \
    "$EMBEDDINGS_URL/discogs-effnet-bs64-1.pb"

# Download classification heads
echo ""
echo "Downloading classification heads..."

# Genre: Discogs400 (professional taxonomy with 400 genres)
download_model "genre_discogs400-discogs-effnet-1.pb" \
    "$CLASS_HEADS_URL/genre_discogs400/genre_discogs400-discogs-effnet-1.pb"

download_model "genre_discogs400-discogs-effnet-1.json" \
    "$CLASS_HEADS_URL/genre_discogs400/genre_discogs400-discogs-effnet-1.json"

# Mood & Instrument: MTG-Jamendo
download_model "mtg_jamendo_moodtheme-discogs-effnet-1.pb" \
    "$CLASS_HEADS_URL/mtg_jamendo_moodtheme/mtg_jamendo_moodtheme-discogs-effnet-1.pb"

download_model "mtg_jamendo_instrument-discogs-effnet-1.pb" \
    "$CLASS_HEADS_URL/mtg_jamendo_instrument/mtg_jamendo_instrument-discogs-effnet-1.pb"

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"