#!/usr/bin/env bash
# ── NAM Model & Cabinet IR Downloader ─────────────────────────────────
#
# Downloads real NAM amp captures (WaveNet/architecture=2) from the
# pelennor2170/NAM_models community collection on GitHub, plus free
# cabinet IR packs from Archive.org.
#
# Models are organized by genre to match the factory presets:
#   Bank 0 — Clean & Edge  (Clean Jazz, Edge of Breakup, Clean w/ Chorus, Dynamic Fingerstyle)
#   Bank 1 — Drive & Lead  (Classic Rock, Lead Solo, Metal Rhythm, Ambient Wash)
#   Bank 2 — FX Bench      (reference presets)
#
# Usage:
#   ./scripts/download_nam_and_irs.sh           # Download everything
#   ./scripts/download_nam_and_irs.sh --list     # List what would be downloaded
#   ./scripts/download_nam_and_irs.sh --models   # Models only
#   ./scripts/download_nam_and_irs.sh --irs      # IRs only
#
# Environment:
#   NAM_DIR:   target directory for .nam files (default: ~/.pedal/models)
#   IR_DIR:    target directory for .wav IRs    (default: ~/.pedal/irs)
#
# ──────────────────────────────────────────────────────────────────────

set -euo pipefail

NAM_DIR="${NAM_DIR:-$HOME/.pedal/models}"
IR_DIR="${IR_DIR:-$HOME/.pedal/irs}"
TMPDIR="${TMPDIR:-/tmp}/pedal-download-$$"
BASE_URL="https://raw.githubusercontent.com/pelennor2170/NAM_models/main"

# ── Colours ──────────────────────────────────────────────────────────
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
RED='\033[0;31m'
CYAN='\033[0;36m'
NC='\033[0m'

info()  { echo -e "${CYAN}[INFO]${NC}  $*"; }
ok()    { echo -e "${GREEN}[OK]${NC}    $*"; }
warn()  { echo -e "${YELLOW}[WARN]${NC}  $*"; }
err()   { echo -e "${RED}[ERR]${NC}   $*" >&2; }

# ── NAM Models by Genre ──────────────────────────────────────────────
#
# Format: "local_filename|github_filename|description|architecture|size_kb"
# All are WaveNet models from the community capture collection.
#
# GENRE MAP (matching factory presets):
#   Clean Jazz, Dynamic Fingerstyle → Fender Twin, Vox AC15 cleans
#   Edge of Breakup → Magnatone pushed, JCM2000 crunch
#   Clean with Chorus → Fender Twin clean
#   Classic Rock → Orange Rockerverb, JCM2000, JCM900
#   Lead Solo → 5150, Engl Savage
#   Metal Rhythm → 5150, Splawn
#   Ambient Wash → Fender Twin clean (with reverb/mod FX)
#

declare -a MODELS=()

# ── Bank 0: Clean & Edge ─────────────────────────────────────────────

MODELS+=(
    "Fender_TwinVerb_Clean.nam|Tim R Fender TwinVerb Norm Bright.nam|Fender Twin Reverb Normal Bright channel — clean jazz/fingerstyle|WaveNet|400"
    "Fender_TwinVerb_Vibrato.nam|Tim R Fender TwinVerb Vibrato Bright.nam|Fender Twin Reverb Vibrato Bright channel — sparkling clean|WaveNet|400"
    "Vox_AC15_TopBoost.nam|Phillipe P VOXAC15-TopBoost.nam|Vox AC15 Top Boost channel — chimey edge-of-breakup|WaveNet|400"
    "Magnatone_Super59_Pushed.nam|Tim R Magnatone Super 59 Mkii Pushed.nam|Magnatone Super 59 pushed — bluesy breakup|WaveNet|400"
    "JCM2000_Clean.nam|Tim R JCM2000 Clean.nam|Marshall JCM2000 clean — crisp rhythm|WaveNet|400"
)

# ── Bank 1: Drive & Lead ─────────────────────────────────────────────

MODELS+=(
    "Orange_Rockerverb.nam|Tom C Axe FX 2 Orange Rockerverb.nam|Orange Rockerverb — classic rock crunch|WaveNet|400"
    "JCM2000_Crunch.nam|Tim R JCM2000 Crunch.nam|Marshall JCM2000 crunch — hard rock rhythm|WaveNet|400"
    "JCM900_Lead.nam|Tim R JCM90050WDualVerbCHAG6.nam|Marshall JCM900 Dual Verb channel A gain 6 — lead|WaveNet|400"
    "5150_BlockLetter_Boosted.nam|Helga B 5150 BlockLetter - Boosted.nam|EVH 5150 Block Letter boosted — high-gain lead|WaveNet|400"
    "6505_Red_NoBoost.nam|Helga B 6505+ Red ch - NoBoost.nam|Peavey 6505+ Red channel no boost — metal rhythm|WaveNet|400"
    "Engl_Savage.nam|Tom C Axe FX 2 Engl Savage.nam|ENGL Savage 120 — modern high-gain lead|WaveNet|400"
    "Splawn_ProMod_3rd.nam|Tim R Splawn Pro Mod 3rd Gear - G12.nam|Splawn Pro Mod 3rd gear — tight high-gain|WaveNet|400"
    "5152_Lead_TS9.nam|Tim R 5152 Lead + TS9.nam|5152 Lead with Tube Screamer — boosted solo tone|WaveNet|400"
)

# ── Pedal captures (for FX chain flexibility) ────────────────────────
# These work as stand-in for the BOOST/OD/DIST blocks or as
# preamp-into-IR tones.

MODELS+=(
    "Klon_Centaur_High.nam|Keith B Klone_plus_BB_highGain.nam|Klon Centaur-style overdrive high gain — boost/overdrive|WaveNet|400"
    "TubeScreamer_Maxed.nam|Tudor N Ibanez Tube Screamer Mini Tone@10_Level@10_Overdrive@0 ESR_0.0006_lite.nam|Ibanez TS808 Mini maxed — clean boost|WaveNet|200"
    "Mesa_MarkIV_Lead.nam|Roman A LT_MESA_MARKIV_1.nam|Mesa Boogie Mark IV lead — singing sustain|WaveNet|400"
    "Bogner_Ecstasy.nam|Tudor N Ceriatone Molecular 50.nam|Bogner Ecstasy-style — versatile overdrive range|WaveNet|400"
)

# ── Feather models (smallest footprint, best for RPi 4B) ────────────

MODELS+=(
    "BOSS_SD1_Feather.nam|Phillipe P BOSS-SD1-Feather-ESR0,001.nam|BOSS SD-1 overdrive feather model — 93 KB, ultra-low CPU|WaveNet|93"
    "MXR_M77_Feather.nam|Phillipe P MXR-M77-Feather-ESR0,001.nam|MXR M77 Custom Badass OD feather model — 93 KB|WaveNet|93"
    "TS808_Feather.nam|Phillipe P Maxon-OD808-Feather-ESR0,001.nam|Maxon OD808 Tube Screamer feather — 93 KB|WaveNet|93"
    "Precision_Drive_Feather.nam|Phillipe P Precision-Drive-Feather-ESR0,001.nam|Precision Drive overdrive feather — 93 KB|WaveNet|93"
    "Friedman_PowerAmp_Feather.nam|Sascha S FriedmanDSM_PowerAmpEL84_MV10_feather.nam|Friedman DSM power amp feather — 93 KB|WaveNet|93"
)

# ── Bank 2: Blues (Tone3000 captures — clean, edge-of-breakup, pushed) ──
# Downloaded from ToneHunt/Tone3000 Supabase storage. "remote_name" is the
# storage key (not a filename), which the Python script resolves to a
# Tone3000 storage URL. The shell script downloads these with direct
# Tone3000 storage URLs via curl.
# Covers: Fender Princeton, Tweed Deluxe, Twin Reverb, Dumble ODS,
# Two-Rock, Marshall Bluesbreaker.

MODELS+=(
    "Fender_Princeton_Clone_Clean.nam|4f81e2b088149c04|Fender Hand-wired Princeton Clone Clean — blues clean (WaveNet A1, 289 KB)|WaveNet|289"
    "Fender_Princeton_Clone_EOB.nam|68a5f3484e1759ed_a2|Fender Hand-wired Princeton Clone Edge-of-Breakup — blues edge (Slimmable A2, 288 KB)|SlimmableContainer|288"
    "Fender_Tweed_Dlx_Edge.nam|ga5rj9sxlnk|Fender Tweed Deluxe The Edge (Celestion Blue) — pushed clean (Slimmable A2, 288 KB)|SlimmableContainer|288"
    "Fender_Twin_Reverb_Clean_A2.nam|tg81uadutjo_a2|Fender Twin Reverb Clean — sparkling blues clean (Slimmable A2, 290 KB)|SlimmableContainer|290"
    "TwoRock_JTM_Signature.nam|qj3hkgbcwb9|Two-Rock John Mayer Signature Prototype — clean/edge bloom (WaveNet A1, 276 KB)|WaveNet|276"
    "Dumble_ODS_102_Rock_A1.nam|5bac98f136812852|Dumble ODS #102 Clone Rock setting — pushed blues overdrive (WaveNet A1, 87 KB feather)|WaveNet|87"
    "Dumble_ODS_102_Rock_A2.nam|5bac98f136812852_a2|Dumble ODS #102 Clone Rock setting — pushed blues OD (Slimmable A2, 289 KB)|SlimmableContainer|289"
    "Marshall_Bluesbreaker_1962_A2.nam|bbc41ecccdc730fb_a2|Marshall Bluesbreaker 1962 — classic blues rock (Slimmable A2, 289 KB)|SlimmableContainer|289"
    "Marshall_Bluesbreaker_Pedal.nam|31aeae1256b36839|Marshall Bluesbreaker Pedal (Original) Setting 1 — pedal OD (WaveNet A1, 290 KB)|WaveNet|290"
)

# ── Total: 31 models (22 core + 9 blues) ──────────────────────────

# ── IR Packs ─────────────────────────────────────────────────────────
# Free cabinet IR collections from Archive.org (CC0 / free for personal use)

declare -a IR_PACKS=(
    "God's Cab|https://archive.org/download/godscab-ir-pack/GodsCab_IR_Pack.zip|godscab.zip|100+ cabinet IRs (412, 1960, V30, Greenback, etc.) — CC0 / Public Domain"
    "Seacow Cabs|https://archive.org/download/seacow-cabs-free-ir-2023/Seacow_Cabs_Free_IRs_2023.zip|seacow.zip|Boutique cab IR pack, 15+ cabs — Free for personal use"
)

# ── Helpers ──────────────────────────────────────────────────────────

_list_models() {
    echo -e "\n${CYAN}NAM Models by Genre:${NC}"
    echo -e "${CYAN}━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━${NC}"
    echo ""

    local current_bank=""
    for entry in "${MODELS[@]}"; do
        IFS='|' read -r local_name remote_name desc arch size <<< "$entry"
        printf "  %-35s %s\n" "$local_name" "$desc"
        printf "  %-35s %s\n" "" "($arch, ~${size} KB)"
        echo ""
    done

    echo -e "${CYAN}Total: ${#MODELS[@]} models${NC}"
}

_download_model() {
    local local_name="$1" remote_name="$2" desc="$3" arch="$4" size="$5"
    local outfile="$NAM_DIR/$local_name"

    if [[ -f "$outfile" ]]; then
        local existing_size
        existing_size=$(stat -c%s "$outfile" 2>/dev/null || echo 0)
        if [[ $existing_size -gt 1000 ]]; then
            ok "SKIP $local_name (already exists, $(( existing_size / 1024 )) KB)"
            return 0
        fi
    fi

    info "DL $local_name ($desc)"

    # URL-encode the remote filename for GitHub raw URL
    local encoded_name
    encoded_name=$(python3 -c "import urllib.parse; print(urllib.parse.quote('$remote_name'))" 2>/dev/null || echo "$remote_name")
    local url="$BASE_URL/$encoded_name"

    local http_code
    http_code=$(curl -sL -o "$outfile" -w "%{http_code}" --connect-timeout 10 --max-time 60 "$url" 2>/dev/null || echo "000")

    if [[ "$http_code" == "200" ]]; then
        local actual_size
        actual_size=$(stat -c%s "$outfile" 2>/dev/null || echo 0)
        if [[ $actual_size -lt 1000 ]]; then
            err "Downloaded file too small (${actual_size} bytes) — removing"
            rm -f "$outfile"
            return 1
        fi
        # Verify it's valid JSON / .nam
        if python3 -c "import json; json.load(open('$outfile'))" 2>/dev/null; then
            ok "Downloaded ${local_name} ($(( actual_size / 1024 )) KB)"
            return 0
        else
            err "Invalid .nam file (not valid JSON) — removing"
            rm -f "$outfile"
            return 1
        fi
    else
        err "HTTP $http_code for $url"
        rm -f "$outfile"
        return 1
    fi
}

_download_all_models() {
    echo -e "\n${CYAN}Downloading ${#MODELS[@]} NAM models to $NAM_DIR${NC}"
    echo -e "${CYAN}━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━${NC}"
    mkdir -p "$NAM_DIR"

    local success=0 fail=0
    for entry in "${MODELS[@]}"; do
        IFS='|' read -r local_name remote_name desc arch size <<< "$entry"
        if _download_model "$local_name" "$remote_name" "$desc" "$arch" "$size"; then
            ((success++))
        else
            ((fail++))
        fi
    done

    echo ""
    info "Models: $success downloaded, $fail failed"
    return "$fail"
}

_download_ir_pack() {
    local name="$1" url="$2" fname="$3" desc="$4"
    local dest="${IR_DIR}/${fname}"

    if [[ -d "$IR_DIR/GodsCab_IR_Pack" ]] || [[ -d "$IR_DIR/Seacow_Cabs_Free_IRs_2023" ]]; then
        local wav_count
        wav_count=$(find "$IR_DIR" -maxdepth 2 -name '*.wav' 2>/dev/null | wc -l)
        if [[ $wav_count -gt 10 ]]; then
            ok "SKIP $name (IRs already extracted, $wav_count .wav files present)"
            return 0
        fi
    fi

    info "Downloading $name..."
    info "  Source: $url"

    if curl -sL -o "$dest" --connect-timeout 10 --max-time 120 "$url" 2>/dev/null; then
        local size
        size=$(stat -c%s "$dest" 2>/dev/null || echo 0)
        if [[ $size -lt 10000 ]]; then
            err "Downloaded file too small (${size} bytes)"
            rm -f "$dest"
            return 1
        fi
        ok "Downloaded ${fname} ($(( size / 1024 / 1024 )) MB)"

        # Extract ZIP
        if command -v unzip &>/dev/null; then
            info "Extracting ${fname} to ${IR_DIR}..."
            unzip -q -o "$dest" -d "$IR_DIR" 2>/dev/null && {
                rm -f "$dest"
                ok "Extracted to ${IR_DIR}/"
                local wav_count
                wav_count=$(find "$IR_DIR" -maxdepth 2 -name '*.wav' 2>/dev/null | wc -l)
                info "IR files installed: $wav_count"
            } || {
                warn "Extraction had issues — ZIP saved at ${dest}"
            }
        else
            warn "unzip not found — ZIP saved at ${dest}"
        fi
        return 0
    else
        err "Failed to download $name from $url"
        return 1
    fi
}

_download_all_irs() {
    echo -e "\n${CYAN}Downloading IR packs to $IR_DIR${NC}"
    echo -e "${CYAN}━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━${NC}"
    mkdir -p "$IR_DIR"

    local success=0 fail=0
    for i in "${!IR_PACKS[@]}"; do
        IFS='|' read -r name url fname desc <<< "${IR_PACKS[$i]}"
        if _download_ir_pack "$name" "$url" "$fname" "$desc"; then
            ((success++))
        else
            ((fail++))
        fi
    done

    echo ""
    info "IR packs: $success downloaded, $fail failed"

    # Count installed IRs
    local wav_count
    wav_count=$(find "$IR_DIR" -maxdepth 2 -name '*.wav' 2>/dev/null | wc -l)
    info "Total .wav IR files in $IR_DIR: $wav_count"
}

_show_counts() {
    echo ""
    local model_count nam_count ir_count
    model_count=$(find "$NAM_DIR" -maxdepth 1 -name '*.nam' 2>/dev/null | wc -l)
    nam_count=$(find "$HOME/.pedal/nam" -maxdepth 1 -name '*.nam' 2>/dev/null | wc -l)
    ir_count=$(find "$IR_DIR" -maxdepth 2 -name '*.wav' 2>/dev/null | wc -l)

    info "Models in $NAM_DIR: $model_count .nam files"
    info "Models in ~/.pedal/nam: $nam_count .nam files (legacy)"
    info "IRs in $IR_DIR: $ir_count .wav files"
}

# ── Main ─────────────────────────────────────────────────────────────

main() {
    mkdir -p "$NAM_DIR" "$IR_DIR"

    case "${1:-}" in
        --list|-l)
            _list_models
            _show_counts
            ;;
        --models|-m)
            _download_all_models
            _show_counts
            ;;
        --irs|-i)
            _download_all_irs
            _show_counts
            ;;
        ""|--all|-a)
            _download_all_models
            echo ""
            _download_all_irs
            echo ""
            _show_counts

            echo ""
            echo -e "${GREEN}╔══════════════════════════════════════════════════════╗${NC}"
            echo -e "${GREEN}║  Download complete!                                 ║${NC}"
            echo -e "${GREEN}╚══════════════════════════════════════════════════════╝${NC}"
            echo ""
            echo "  Next steps:"
            echo "  1. Run the inventory report:"
            echo "     python3 scripts/models/nam/build_inventory.py"
            echo ""
            echo "  2. Verify models load correctly:"
            echo "     python3 -c \"from src.dsp.nam_host import NAMHost;"
            echo "       h = NAMHost(models_dir='$NAM_DIR');"
            echo "       for m in h.list_available_models(): print(m.name, m.architecture, f'{m.size_mb:.1f}MB')\""
            echo ""
            echo "  3. Check IRs:"
            echo "     python3 -c \"from pathlib import Path;"
            echo "       irs = list(Path('$IR_DIR').rglob('*.wav'));"
            echo "       print(f'{len(irs)} IRs loaded')\""
            echo ""
            ;;
        --help|-h)
            echo "Usage: $0 [OPTION]"
            echo ""
            echo "Options:"
            echo "  --list, -l     List models and IRs (without downloading)"
            echo "  --models, -m   Download NAM models only"
            echo "  --irs, -i      Download IR packs only"
            echo "  --all, -a      Download everything (default)"
            echo "  --help, -h     Show this help"
            echo ""
            echo "Environment:"
            echo "  NAM_DIR   target for .nam files (default: ~/.pedal/models)"
            echo "  IR_DIR    target for .wav IRs   (default: ~/.pedal/irs)"
            ;;
        *)
            err "Unknown option: $1"
            echo "Usage: $0 [--list|--models|--irs|--all|--help]"
            exit 1
            ;;
    esac
}

main "$@"