#!/usr/bin/env python3
"""Download NAM WaveNet models from community collection + cabinet IRs.

Downloads real NAM amp captures from the pelennor2170/NAM_models GitHub
repository and free cabinet IR packs from Archive.org.

Usage:
    python3 scripts/download_nam_and_irs.py          # Download everything
    python3 scripts/download_nam_and_irs.py --list   # List models only
    python3 scripts/download_nam_and_irs.py --models # Models only
    python3 scripts/download_nam_and_irs.py --irs    # IRs only
"""

import json
import os
import subprocess
import sys
import urllib.parse
import urllib.request
from pathlib import Path

# ── Paths ────────────────────────────────────────────────────────────────
NAM_DIR = Path.home() / ".pedal" / "models"
IR_DIR = Path.home() / ".pedal" / "irs"
BASE_URL = "https://raw.githubusercontent.com/pelennor2170/NAM_models/main"

# ── Model definitions ────────────────────────────────────────────────────
# (local_name, remote_name, description, architecture, size_kb)
# URL for Tone3000 (ToneHunt) Supabase storage downloads
TONE3000_STORAGE = "https://api.tone3000.com/storage/v1/object/public/models"

MODELS = [
    # ── Bank 0: Clean & Edge ──────────────────────────────────────────
    ("Fender_TwinVerb_Clean.nam",
     "Tim R Fender TwinVerb Norm Bright.nam",
     "Fender Twin Reverb Normal Bright — clean jazz/fingerstyle",
     "WaveNet", 400),
    ("Fender_TwinVerb_Vibrato.nam",
     "Tim R Fender TwinVerb Vibrato Bright.nam",
     "Fender Twin Reverb Vibrato Bright — sparkling clean",
     "WaveNet", 400),
    ("Vox_AC15_TopBoost.nam",
     "Phillipe P VOXAC15-TopBoost.nam",
     "Vox AC15 Top Boost — 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 ──────────────────────────────────────────
    ("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 ch.A gain6 — 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 ch 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",
     "WaveNet", 400),

    # ── Pedal captures ───────────────────────────────────────────────
    ("Klon_Centaur_High.nam",
     "Keith B Klone_plus_BB_highGain.nam",
     "Klon Centaur-style OD 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 OD range",
     "WaveNet", 400),

    # ── Feather models (93 KB, ultra-low CPU) ────────────────────────
    ("BOSS_SD1_Feather.nam",
     "Phillipe P BOSS-SD1-Feather-ESR0,001.nam",
     "BOSS SD-1 overdrive feather — 93 KB",
     "WaveNet", 93),
    ("MXR_M77_Feather.nam",
     "Phillipe P MXR-M77-Feather-ESR0,001.nam",
     "MXR M77 Custom Badass OD feather — 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. These are real
    # amp captures of classic blues tones: Fender Princeton, Tweed Deluxe,
    # Twin Reverb, Dumble ODS, Two-Rock, Marshall Bluesbreaker.
    # Format: (local_name, tone3000_storage_key, description, architecture, size_kb)
    ("Fender_Princeton_Clone_Clean.nam",
     "4f81e2b088149c04.nam",
     "Fender Hand-wired Princeton Clone Clean — blues clean (WaveNet A1, 289 KB)",
     "WaveNet", 289),
    ("Fender_Princeton_Clone_EOB.nam",
     "68a5f3484e1759ed_a2.nam",
     "Fender Hand-wired Princeton Clone Edge-of-Breakup — blues edge (Slimmable A2, 288 KB)",
     "SlimmableContainer", 288),
    ("Fender_Tweed_Dlx_Edge.nam",
     "ga5rj9sxlnk.nam",
     "Fender Tweed Deluxe The Edge (Celestion Blue) — pushed clean (Slimmable A2, 288 KB)",
     "SlimmableContainer", 288),
    ("Fender_Twin_Reverb_Clean_A2.nam",
     "tg81uadutjo_a2.nam",
     "Fender Twin Reverb Clean — sparkling blues clean (Slimmable A2, 290 KB)",
     "SlimmableContainer", 290),
    ("TwoRock_JTM_Signature.nam",
     "qj3hkgbcwb9.nam",
     "Two-Rock John Mayer Signature Prototype — clean/edge bloom (WaveNet A1, 276 KB)",
     "WaveNet", 276),
    ("Dumble_ODS_102_Rock_A1.nam",
     "5bac98f136812852.nam",
     "Dumble ODS #102 Clone Rock setting — pushed blues overdrive (WaveNet A1, 87 KB feather)",
     "WaveNet", 87),
    ("Dumble_ODS_102_Rock_A2.nam",
     "5bac98f136812852_a2.nam",
     "Dumble ODS #102 Clone Rock setting — pushed blues OD (Slimmable A2, 289 KB)",
     "SlimmableContainer", 289),
    ("Marshall_Bluesbreaker_1962_A2.nam",
     "bbc41ecccdc730fb_a2.nam",
     "Marshall Bluesbreaker 1962 — classic blues rock (Slimmable A2, 289 KB)",
     "SlimmableContainer", 289),
    ("Marshall_Bluesbreaker_Pedal.nam",
     "31aeae1256b36839.nam",
     "Marshall Bluesbreaker Pedal (Original) Setting 1 — pedal OD (WaveNet A1, 290 KB)",
     "WaveNet", 290),


    # ── Bank 7: Bass, Orange & Others (Tone3000 A2 — SlimmableContainer) ────
    # Downloaded from Tone3000 Supabase storage. Covers missing categories:
    # Bass (Ampeg, Aguilar, GK), Orange Rocker 30, Fender Bassman,
    # Hiwatt DR103, Supro Black Magick.
    ("Ampeg_PF20T_A2.nam",
     "db2c43e99cfc346d_a2.nam",
     "Ampeg PF20T — clean bass tube amp (Slimmable A2, 288 KB)",
     "SlimmableContainer", 288),
    ("Aguilar_Tone_Hammer_500_A2.nam",
     "38rdhwzk3hp_a2.nam",
     "Aguilar Tone Hammer 500 — modern bass preamp (Slimmable A2, 293 KB)",
     "SlimmableContainer", 293),
    ("GK_800RB_Line_A2.nam",
     "ce68ccd373a10054_a2.nam",
     "Gallien-Krueger 800RB — classic bass solid-state (Slimmable A2, 295 KB)",
     "SlimmableContainer", 295),
    ("Aguilar_DB751_Cab_A2.nam",
     "lage0x4ybk_a2.nam",
     "Aguilar DB751 — bass amp DI (Slimmable A2, 292 KB)",
     "SlimmableContainer", 292),
    ("Orange_Rocker30_Dirty_A2.nam",
     "8491c70c3771de1e_a2.nam",
     "Orange Rocker 30 Dirty — classic rock crunch (Slimmable A2, 295 KB)",
     "SlimmableContainer", 295),
    ("Orange_Rocker30_Natural_A2.nam",
     "a60ed1a1aa0c2c0f_a2.nam",
     "Orange Rocker 30 Natural — clean edge (Slimmable A2, 289 KB)",
     "SlimmableContainer", 289),
    ("Fender_Bassman_Dimed_A2.nam",
     "0a9223a9f1c879d5_a2.nam",
     "Fender Bassman dimed — blues/rock breakup (Slimmable A2, 289 KB)",
     "SlimmableContainer", 289),
    ("Hiwatt_DR103_Bright_A2.nam",
     "0c46e5a7db1cce48_a2.nam",
     "Hiwatt DR-103 Bright Normal — prog/classic rock (Slimmable A2, 287 KB)",
     "SlimmableContainer", 287),
    ("Hiwatt_DR103_Cornish_A2.nam",
     "ynk0h2e2tph.nam",
     "Hiwatt DR103 + Cornish PGS — boosted prog rock (Slimmable A2, 288 KB)",
     "SlimmableContainer", 288),
    ("Supro_Black_Magick_M3_A2.nam",
     "c90a6d8ee0c0a689_a2.nam",
     "Supro Black Magick M3 — garage/blues (Slimmable A2, 288 KB)",
     "SlimmableContainer", 288),
]

# ── IR Pack definitions ──────────────────────────────────────────────────
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"),
]


def info(msg):    print(f"\033[0;36m[INFO]\033[0m  {msg}")
def ok(msg):      print(f"\033[0;32m[OK]\033[0m    {msg}")
def warn(msg):    print(f"\033[1;33m[WARN]\033[0m  {msg}")
def err(msg):     print(f"\033[0;31m[ERR]\033[0m   {msg}", file=sys.stderr)


def download_file(url: str, dest: Path, max_size_mb: int = 50) -> bool:
    """Download a file from url to dest. Returns True on success."""
    try:
        req = urllib.request.Request(url, headers={"User-Agent": "pedal-downloader/1.0"})
        with urllib.request.urlopen(req, timeout=30) as src:
            data = src.read()
        size_mb = len(data) / (1024 * 1024)
        if size_mb > max_size_mb:
            err(f"File too large ({size_mb:.1f} MB, max {max_size_mb} MB)")
            return False
        if len(data) < 1000:
            err(f"File too small ({len(data)} bytes)")
            return False
        dest.write_bytes(data)
        return True
    except Exception as e:
        err(f"Download failed: {e}")
        return False


def download_model(local_name: str, remote_name: str, desc: str, arch: str, size: int) -> bool:
    """Download a single NAM model file from GitHub."""
    outfile = NAM_DIR / local_name

    if outfile.exists() and outfile.stat().st_size > 1000:
        ok(f"SKIP {local_name} (already exists, {outfile.stat().st_size // 1024} KB)")
        return True

    # URL-encode the remote filename
    encoded = urllib.parse.quote(remote_name)
    url = f"{BASE_URL}/{encoded}"

    info(f"DL {local_name} ({desc})")
    if download_file(url, outfile):
        # Verify it's valid JSON (.nam format)
        try:
            data = json.loads(outfile.read_text())
            actual_arch = data.get("architecture", "?")
            actual_ver = data.get("version", "?")
            ok(f"Downloaded {local_name} ({outfile.stat().st_size // 1024} KB, {actual_arch} v{actual_ver})")
            return True
        except json.JSONDecodeError as e:
            err(f"Invalid .nam file (not valid JSON): {e}")
            outfile.unlink(missing_ok=True)
            return False
    else:
        err(f"Failed to download {local_name}")
        return False


def _discover_tone3000_key() -> str:
    """Discover the Supabase anon key from Tone3000 website JS bundles."""
    try:
        req = urllib.request.Request(
            "https://tone3000.com",
            headers={"User-Agent": "Mozilla/5.0 (compatible; pedal-downloader/1.0)"},
        )
        with urllib.request.urlopen(req, timeout=15) as resp:
            html = resp.read().decode("utf-8", errors="replace")
        import re
        chunk_urls = set()
        for m in re.finditer(r'src="([^"]*\.js[^"]*)"', html):
            url = m.group(1)
            if url.startswith("/"):
                url = "https://tone3000.com" + url
            chunk_urls.add(url)
        for chunk_url in sorted(chunk_urls)[:15]:
            try:
                req2 = urllib.request.Request(chunk_url, headers={"User-Agent": "Mozilla/5.0"})
                with urllib.request.urlopen(req2, timeout=10) as resp2:
                    js = resp2.read().decode("utf-8", errors="replace")
                m = re.search(
                    r'createBrowserClient\s*\)\s*\(\s*["\']https?://api\.tone3000\.com["\']\s*,\s*["\']([^"\']+)["\']',
                    js,
                )
                if m:
                    key = m.group(1)
                    if key.startswith("eyJ") and key.count(".") >= 2:
                        return key
                for m2 in re.finditer(
                    r'["\'](eyJ[a-zA-Z0-9_\-]{20,}\.[a-zA-Z0-9_\-]{20,}\.[a-zA-Z0-9_\-]{10,})["\']', js
                ):
                    key = m2.group(1)
                    if key.count(".") >= 2:
                        return key
            except Exception:
                continue
    except Exception:
        pass
    return ""


def download_tone3000_model(local_name: str, storage_key: str, desc: str, arch: str, size: int) -> bool:
    """Download a single NAM model file from Tone3000 Supabase storage."""
    outfile = NAM_DIR / local_name

    if outfile.exists() and outfile.stat().st_size > 1000:
        ok(f"SKIP {local_name} (already exists, {outfile.stat().st_size // 1024} KB)")
        return True

    url = f"{TONE3000_STORAGE}/{storage_key}"

    info(f"DL {local_name} ({desc})")
    req = urllib.request.Request(url, headers={"User-Agent": "pedal-downloader/1.0"})
    try:
        with urllib.request.urlopen(req, timeout=30) as src:
            data = src.read()
    except Exception as e:
        err(f"Download failed: {e}")
        # Try with anon key (some storage buckets require authentication)
        try:
            anon_key = _discover_tone3000_key()
            if anon_key:
                req = urllib.request.Request(url, headers={
                    "User-Agent": "pedal-downloader/1.0",
                    "Authorization": f"Bearer {anon_key}",
                })
                with urllib.request.urlopen(req, timeout=30) as src:
                    data = src.read()
            else:
                return False
        except Exception as e2:
            err(f"Download failed even with auth: {e2}")
            return False

    size_mb = len(data) / (1024 * 1024)
    if size_mb > 50:
        err(f"File too large ({size_mb:.1f} MB, max 50 MB)")
        return False
    if len(data) < 1000:
        err(f"File too small ({len(data)} bytes)")
        return False
    outfile.write_bytes(data)

    # Verify
    try:
        model_data = json.loads(data)
        actual_arch = model_data.get("architecture", "?")
        actual_ver = model_data.get("version", "?")
        ok(f"Downloaded {local_name} ({outfile.stat().st_size // 1024} KB, {actual_arch} v{actual_ver})")
        return True
    except json.JSONDecodeError as e:
        err(f"Invalid .nam file (not valid JSON): {e}")
        outfile.unlink(missing_ok=True)
        return False


def download_models():
    """Download all NAM models (GitHub + Tone3000)."""
    print(f"\n\033[0;36mDownloading {len(MODELS)} NAM models to {NAM_DIR}\033[0m")
    print(f"\033[0;36m{'━' * 60}\033[0m")
    NAM_DIR.mkdir(parents=True, exist_ok=True)

    success = 0
    fail = 0

    for local_name, remote_name, desc, arch, size in MODELS:
        # Detect Tone3000 storage key (short, no dots, no spaces)
        if " " not in remote_name and "." not in remote_name and len(remote_name) < 64:
            if download_tone3000_model(local_name, remote_name, desc, arch, size):
                success += 1
            else:
                fail += 1
        else:
            if download_model(local_name, remote_name, desc, arch, size):
                success += 1
            else:
                fail += 1

    print(f"\n  Models: {success} successful, {fail} failed")
    return fail


def download_irs():
    """Download IR packs from Archive.org."""
    print(f"\n\033[0;36mDownloading IR packs to {IR_DIR}\033[0m")
    print(f"\033[0;36m{'━' * 60}\033[0m")
    IR_DIR.mkdir(parents=True, exist_ok=True)

    success = 0
    fail = 0

    for name, url, fname, desc in IR_PACKS:
        # Check if already extracted
        wav_count = len(list(IR_DIR.rglob("*.wav")))
        if wav_count > 10:
            ok(f"SKIP {name} ({wav_count} .wav files already present)")
            success += 1
            continue

        dest = IR_DIR / fname
        info(f"Downloading {name}...")
        info(f"  Source: {url}")

        if download_file(url, dest, max_size_mb=200):
            size_mb = dest.stat().st_size / (1024 * 1024)
            ok(f"Downloaded {fname} ({size_mb:.1f} MB)")

            # Extract ZIP
            import zipfile
            try:
                with zipfile.ZipFile(dest, "r") as zf:
                    zf.extractall(IR_DIR)
                dest.unlink()
                wav_count = len(list(IR_DIR.rglob("*.wav")))
                ok(f"Extracted to {IR_DIR}/ ({wav_count} .wav files)")
                success += 1
            except zipfile.BadZipFile:
                warn(f"Bad ZIP — saved at {dest}")
                fail += 1
        else:
            err(f"Failed to download {name}")
            fail += 1

    print(f"\n  IR packs: {success} downloaded, {fail} failed")
    wav_count = len(list(IR_DIR.rglob("*.wav")))
    info(f"Total .wav IR files: {wav_count}")
    return fail


def list_models():
    """Print model listing."""
    print(f"\n\033[0;36mNAM Models by Genre:\033[0m")
    print(f"\033[0;36m{'━' * 60}\033[0m")
    for local_name, remote_name, desc, arch, size in MODELS:
        print(f"  {local_name:35s} {desc}")
        print(f"  {'':35s} ({arch}, ~{size} KB)")
        print()

    print(f"\nTotal: {len(MODELS)} models")

    # Count existing files
    nam_count = len(list(NAM_DIR.glob("*.nam")))
    ir_count = len(list(IR_DIR.rglob("*.wav")))
    print(f"\n  Models in {NAM_DIR}: {nam_count}")
    print(f"  IRs in {IR_DIR}:      {ir_count}")


def main():
    # Also check ~/.pedal/nam/ legacy location
    NAM_DIR.mkdir(parents=True, exist_ok=True)
    IR_DIR.mkdir(parents=True, exist_ok=True)

    arg = sys.argv[1] if len(sys.argv) > 1 else "--all"

    if arg == "--list":
        list_models()
    elif arg == "--models":
        exit(download_models())
    elif arg == "--irs":
        exit(download_irs())
    elif arg == "--all":
        fail = download_models()
        download_irs()

        # Summary
        nam_count = len(list(NAM_DIR.glob("*.nam")))
        ir_count = len(list(IR_DIR.rglob("*.wav")))
        print(f"\n\033[0;32m{'╔' + '═' * 56 + '╗'}\033[0m")
        print(f"\033[0;32m║  Download complete!{' ' * 36}║\033[0m")
        print(f"\033[0;32m{'╚' + '═' * 56 + '╝'}\033[0m")
        print()
        print(f"  Models: {nam_count} .nam files in {NAM_DIR}")
        print(f"  IRs:    {ir_count} .wav files in {IR_DIR}")
        exit(fail)
    else:
        print(f"Usage: {sys.argv[0]} [--list|--models|--irs|--all]")
        exit(1)


if __name__ == "__main__":
    main()
