#!/bin/bash

# --- CONFIGURATION ---
# 1. The Main (COMPLETE) Dataset
# Composed of 10 images, 4 cubes from them, 6 derivatives from each cube.
MAIN_DATASET="https://www.kaggle.com/api/v4/datasets/download/cristianllull/shrec-2026-retrieval-of-high-frequency-geometry"

# 2. Sample Dataset
SAMPLE_DATASET="https://www.kaggle.com/api/v1/datasets/download/cristianllull/shrec-2026-sample-placeholder"

# --- AUTH CHECK ---
if [[ -z "$KAGGLE_USERNAME" ]] || [[ -z "$KAGGLE_KEY" ]]; then
    echo "Error: Credentials not found."
    echo "Please export your variables:"
    echo "  export KAGGLE_USERNAME=your_username"
    echo "  export KAGGLE_KEY=your_key"
    exit 1
fi

echo "--- SHREC 2026 Downloader ---"
echo "Logged in as: $KAGGLE_USERNAME"

download_dataset() {
    local dataset_dir=$1
    local output_file=$2
    
    echo "Target: $dataset_dir"
    echo "Saving to: $output_file"
    echo "Starting download (Resumable)..."

    # -C - allows resuming if the internet cuts out
    # -L follows redirects
    curl -L -C - -u "$KAGGLE_USERNAME:$KAGGLE_KEY" \
        -o "$output_file" \
        "$dataset_dir"

    # Verify it is not an error HTML page
    if file "$output_file" | grep -q "HTML"; then
        echo "CRITICAL ERROR: The download link returned an HTML page instead of a ZIP."
        echo "This usually means the Dataset ID is wrong or your Credentials are invalid."
        rm "$output_file"
        exit 1
    fi

    echo "Download Complete: $output_file"
    echo "File size: $(du -h "$output_file" | cut -f1)"
}

# --- MAIN LOGIC ---

case "$1" in
    --download_sample)
        echo "Mode: SAMPLE DATA (Lightweight)"
        download_dataset "$SAMPLE_DATASET" "Data_Sample.zip"
        
        # Safe to auto-extract the sample because it is small
        echo "Extracting sample..."
        unzip -q -o "Data_Sample.zip"
        echo "Done. Check 'Data_Sample/' folder."
        ;;

    --download_all)
        echo "Mode: FULL DATASET (17 GB)"
        echo "WARNING: This requires stable internet."
        download_dataset "$MAIN_DATASET" "SHREC_2026_Full.zip"
        
        echo ""
        echo "---------------------------------------------------"
        echo "DOWNLOAD SUCCESSFUL."
        echo "The file 'SHREC_2026_Full.zip' is saved."
        echo "Unzipping this file requires ~50GB of free disk space."
        echo "To unzip manually, run: unzip SHREC_2026_Full.zip"
        echo "---------------------------------------------------"
        ;;

    *)
        echo "Usage: $0 {--download_sample | --download_all}"
        exit 1
        ;;
esac
