"""
Recategorize ALL samples in sample_metadata.db with clean, normalized categories.

Maps the messy folder-based categories (e.g. "LATINOS - DRUM LOOPS", "33 Instrumental Loops")
to clean pipeline-ready categories: kick, snare, hihat, clap, drumloops, bass, perc,
fx, impact, synth, keys, pad, vocals, oneshots, melody, chords, guitar, brass, bells, fill.

Also adds MIDI files from SentimientoLatino2025 and reggaeton 3 to the DB.

Usage:
    python recategorize_samples.py
"""

import os
import sys
import sqlite3
from pathlib import Path

DB_PATH = Path(r"C:\ProgramData\Ableton\Live 12 Suite\Resources\MIDI Remote Scripts\libreria\reggaeton\sample_metadata.db")
LIBRERIA = Path(r"C:\ProgramData\Ableton\Live 12 Suite\Resources\MIDI Remote Scripts\libreria\reggaeton")

CLEAN_CATEGORIES = {
    "kick", "snare", "hihat", "clap", "drumloops", "bass", "perc",
    "fx", "impact", "synth", "keys", "pad", "vocals", "oneshots",
    "melody", "chords", "guitar", "brass", "bells", "fill", "music_loop",
}


def _infer_clean_category(rel_path: str, filename: str) -> str:
    """Infer a clean category from path and filename.

    Priority: filename keywords > path keywords > folder name.
    """
    path_lower = rel_path.lower().replace("\\", "/")
    fn_lower = filename.lower()

    # --- Filename-based detection (highest priority) ---

    # Drum hits (oneshots)
    if "kick" in fn_lower and "loop" not in fn_lower:
        return "kick"
    if "snare" in fn_lower and "loop" not in fn_lower:
        return "snare"
    if any(kw in fn_lower for kw in ("hi-hat", "hihat", "hi hat", "hh")):
        if "loop" in fn_lower:
            return "drumloops"
        return "hihat"
    if "clap" in fn_lower and "loop" not in fn_lower:
        return "clap"
    if "rim" in fn_lower and "loop" not in fn_lower:
        return "perc"

    # Bass
    if any(kw in fn_lower for kw in ("bass", "sub bass", "sub_", "reese", "resse", "808")):
        if "loop" in fn_lower or "music" in path_lower:
            return "music_loop"
        return "bass"

    # FX and impacts
    if any(kw in fn_lower for kw in ("impact", "camtazo", "hit")):
        return "impact"
    if any(kw in fn_lower for kw in ("riser", "sweep", "transition", "fx", "fx_")):
        return "fx"
    if "fill" in fn_lower:
        return "fill"

    # Percussion loops / fills
    if "perc" in fn_lower and ("loop" in fn_lower or path_lower.count("/") <= 2):
        return "perc"

    # --- Path-based detection ---

    # reggaeton 3 specific folders
    if "reggaeton 3" in path_lower:
        if "/8. kicks" in path_lower:
            return "kick"
        if "/9. snare" in path_lower:
            return "snare"
        if "/10. percs" in path_lower:
            return "perc"
        if "/4. drum loops" in path_lower:
            return "drumloops"
        if "/5. fx" in path_lower:
            return "fx"
        if "/6. impact" in path_lower:
            return "impact"
        if "/7. fill" in path_lower:
            return "fill"
        if "/11. vocals" in path_lower:
            return "vocals"
        if "/3. one shots" in path_lower:
            return "oneshots"

    # SentimientoLatino2025 /01/ specific folders
    if "sentimientolatino2025" in path_lower:
        if "drum loops" in path_lower:
            return "drumloops"
        if "one shots" in path_lower:
            return "oneshots"
        if "midi pack" in path_lower:
            return "chords"

    # SentimientoLatino2025 /02/ specific folders
    if "/02/" in path_lower and "sentimientolatino2025" in path_lower:
        if "drum loops" in path_lower or "/23 " in path_lower:
            return "drumloops"
        if "music loops" in path_lower or "/07 " in path_lower:
            return "music_loop"
        if "instrumental loops" in path_lower or "/33 " in path_lower:
            # Instrumental loops contain bass, keys, pads, etc
            pass  # fall through to filename analysis
        if "one shots" in path_lower or "/20 " in path_lower:
            return "oneshots"
        if "vocals" in path_lower:
            return "vocals"

    # --- Filename keyword-based for sample pack subfolders ---

    # Drum loops (filename patterns)
    if "loop" in fn_lower and "drum" in path_lower:
        return "drumloops"
    if "loop" in fn_lower and "perc" in path_lower:
        return "perc"
    if any(kw in fn_lower for kw in ("drumloop", "drum_loop")):
        return "drumloops"

    # SentimientoLatino2025 sample pack items - detect from filename keywords
    if "_drums" in fn_lower:
        return "drumloops"
    if "_drum" in fn_lower:
        return "drumloops"
    if "_perc" in fn_lower:
        return "perc"
    if "_snare" in fn_lower:
        return "snare"

    # Instruments
    if any(kw in fn_lower for kw in ("chord", "bell_chord")):
        return "chords"
    if any(kw in fn_lower for kw in ("pad", "texture")):
        return "pad"
    if any(kw in fn_lower for kw in ("lead", "pluck")):
        return "synth"
    if any(kw in fn_lower for kw in ("arp", "arpeggio")):
        return "synth"
    if any(kw in fn_lower for kw in ("rhode", "rhodes", "piano", "keys")):
        return "keys"
    if any(kw in fn_lower for kw in ("guitar",)):
        return "guitar"
    if any(kw in fn_lower for kw in ("vocal", "vox", "voice")):
        return "vocals"
    if any(kw in fn_lower for kw in ("brass",)):
        return "brass"
    if any(kw in fn_lower for kw in ("bell", "mallet")):
        return "bells"
    if any(kw in fn_lower for kw in ("synth",)):
        return "synth"
    if any(kw in fn_lower for kw in ("cymatics", "fx", "transition", "riser")):
        return "fx"

    # Main libreria folders
    if path_lower.startswith("kick/"):
        return "kick"
    if path_lower.startswith("snare/"):
        return "snare"
    if "hi-hat" in path_lower or "hihat" in path_lower:
        return "hihat"
    if path_lower.startswith("drumloops/"):
        return "drumloops"
    if path_lower.startswith("perc loop/"):
        return "perc"
    if path_lower.startswith("bass/"):
        return "bass"
    if path_lower.startswith("fx/"):
        return "fx"
    if path_lower.startswith("oneshots/"):
        return "oneshots"

    # Music loops
    if "music" in path_lower and "loop" in path_lower:
        return "music_loop"

    # Vocals
    if "vocal" in path_lower:
        return "vocals"

    # Instrumental loops - categorize by content
    if "instrumental" in path_lower:
        if "bass" in fn_lower:
            return "bass"
        if "pad" in fn_lower:
            return "pad"
        if "keys" in fn_lower:
            return "keys"
        if "fx" in fn_lower or "vocal" in fn_lower or "chop" in fn_lower:
            return "fx"
        return "synth"

    return "other"


def recategorize():
    conn = sqlite3.connect(str(DB_PATH))
    conn.row_factory = sqlite3.Row
    c = conn.cursor()

    c.execute("SELECT path FROM samples")
    rows = c.fetchall()

    # Clear all old categories
    c.execute("DELETE FROM sample_categories")

    updated = 0
    category_counts = {}

    for row in rows:
        path = row["path"]
        filename = os.path.basename(path)

        category = _infer_clean_category(path, filename)

        c.execute(
            "INSERT OR IGNORE INTO sample_categories (path, category) VALUES (?, ?)",
            (path, category)
        )

        category_counts[category] = category_counts.get(category, 0) + 1
        updated += 1

    conn.commit()

    print(f"Recategorized {updated} samples")
    print("\nCategory distribution:")
    for cat, count in sorted(category_counts.items(), key=lambda x: -x[1]):
        print(f"  {cat:15s}: {count:4d}")

    conn.close()
    return category_counts


def add_midi_files():
    """Add MIDI files from SentimientoLatino2025 and reggaeton 3 to DB."""
    conn = sqlite3.connect(str(DB_PATH))
    c = conn.cursor()

    added = 0
    for root, dirs, files in os.walk(str(LIBRERIA)):
        for f in files:
            if not f.lower().endswith(('.mid', '.midi')):
                continue

            full_path = os.path.join(root, f)
            rel_path = os.path.relpath(full_path, str(LIBRERIA))

            c.execute("SELECT 1 FROM samples WHERE path = ?", (rel_path,))
            if c.fetchone():
                continue

            # Parse BPM and key from MIDI filename
            sys.path.insert(0, str(Path(__file__).parent))
            from bpm_key_parser import parse_bpm, parse_key
            parsed_bpm = parse_bpm(f)
            parsed_key = parse_key(f)

            c.execute(
                "INSERT OR IGNORE INTO samples (path, bpm, key, analyzed_at) VALUES (?, ?, ?, datetime('now'))",
                (rel_path, parsed_bpm, parsed_key)
            )

            # Infer category for MIDI
            fn_lower = f.lower()
            if "chord" in fn_lower or "progres" in fn_lower:
                cat = "chords"
            elif "arp" in fn_lower:
                cat = "synth"
            elif "bass" in fn_lower:
                cat = "bass"
            elif "drum" in fn_lower:
                cat = "drumloops"
            elif "lead" in fn_lower:
                cat = "synth"
            elif "melody" in fn_lower:
                cat = "melody"
            elif "pad" in fn_lower:
                cat = "pad"
            elif "piano" in fn_lower or "rhode" in fn_lower:
                cat = "keys"
            else:
                cat = "chords"

            c.execute(
                "INSERT OR IGNORE INTO sample_categories (path, category) VALUES (?, ?)",
                (rel_path, cat)
            )

            added += 1

    conn.commit()
    conn.close()
    print(f"Added {added} MIDI files to DB")


def verify():
    conn = sqlite3.connect(str(DB_PATH))
    c = conn.cursor()

    c.execute("SELECT COUNT(*) FROM samples")
    total = c.fetchone()[0]
    c.execute("SELECT COUNT(DISTINCT path) FROM sample_categories")
    categorized = c.fetchone()[0]
    c.execute("SELECT COUNT(*) FROM samples WHERE bpm > 0")
    with_bpm = c.fetchone()[0]
    c.execute("SELECT COUNT(*) FROM samples WHERE key IS NOT NULL AND key != ''")
    with_key = c.fetchone()[0]

    print(f"\n{'='*50}")
    print(f"DB Verification")
    print(f"{'='*50}")
    print(f"Total samples:        {total}")
    print(f"With categories:      {categorized}")
    print(f"With BPM > 0:         {with_bpm}")
    print(f"With key:             {with_key}")

    # Show samples per source
    c.execute("""
        SELECT 
            CASE 
                WHEN path LIKE 'SentimientoLatino%' THEN 'SentimientoLatino2025'
                WHEN path LIKE 'reggaeton 3%' THEN 'reggaeton 3'
                ELSE 'main library'
            END as source,
            COUNT(*) as count
        FROM samples
        GROUP BY source
    """)
    print("\nBy source:")
    for row in c.fetchall():
        print(f"  {row[0]:30s}: {row[1]:4d}")

    # Show category distribution by source
    for source, pattern in [
        ("SentimientoLatino2025", "SentimientoLatino%"),
        ("reggaeton 3", "reggaeton 3%"),
        ("main library", "kick/%"),
    ]:
        if source == "main library":
            print(f"\n{'-- main library categories --'}")
            c.execute("""
                SELECT sc.category, COUNT(*) 
                FROM sample_categories sc
                JOIN samples s ON sc.path = s.path
                WHERE s.path NOT LIKE 'SentimientoLatino%' AND s.path NOT LIKE 'reggaeton 3%'
                GROUP BY sc.category ORDER BY COUNT(*) DESC
            """)
        else:
            print(f"\n{'-- ' + source + ' categories --'}")
            c.execute("""
                SELECT sc.category, COUNT(*) 
                FROM sample_categories sc
                JOIN samples s ON sc.path = s.path
                WHERE s.path LIKE ?
                GROUP BY sc.category ORDER BY COUNT(*) DESC
            """, (pattern,))
        for row in c.fetchall():
            print(f"  {row[0]:15s}: {row[1]:4d}")

    conn.close()


if __name__ == "__main__":
    print("Phase 1: Recategorize all samples with clean categories...")
    recategorize()

    print("\nPhase 2: Add MIDI files to DB...")
    add_midi_files()

    print("\nPhase 3: Verify...")
    verify()

    print("\nDone!")
