#!/usr/bin/env python3
"""
clear-bad-genres.py — normalize or blank GENRE tags so every file ends up
with a clean "; "-joined list of whitelist entries (max 3 parts).

Two-tier behavior per file:
  1. NORMALIZE: if every part of the GENRE value canonicalizes to a
     whitelist entry (after splitting on ; , / or " - " and resolving aliases
     like "Rnb" → "R&B", "Hip Hop" → "Hip-Hop"), and there are ≤3 parts,
     REWRITE the file's GENRE in canonical "Display1; Display2" form. This
     fixes formatting issues like "Dance;Pop" (no space) → "Dance; Pop"
     without losing the data.
  2. BLANK: if any part doesn't map to a whitelist canonical, OR there are
     too many parts, BLANK the GENRE field. Run spotify-genre.py --apply
     --force afterwards to refill what Spotify can.

Catches:
  - Mojibake (Cyrillic-as-ASCII like '5:B@>==0O <C7K:0')
  - Wrong delimiters: comma ('Pop, Rock'), slash ('Indie Rock/Rock Pop'),
    space-dash-space ('Dance - Electro - Vocal'), missing space after
    semicolon ('Dance;Pop')
  - Mixed-delimiter chains ('Pop, Rock; Alternative')
  - Bracketed values ('[Drum n Bass]')
  - Over-tagging: more than 3 parts (matches spotify-genre.py's MAX_GENRES)
  - Spelling variants: 'Rnb' → 'R&B', 'Hip Hop' → 'Hip-Hop',
    'Drum n Bass' / 'Drum and Bass' → 'Drum & Bass', 'K Pop' → 'K-Pop'
  - Niche values not in the whitelist

Already-clean tracks are left untouched (zero file writes).

Usage:
  clear-bad-genres.py             # dry run, show planned normalize+blank ops
  clear-bad-genres.py --apply     # actually rewrite/blank
"""
import sys, os, re, argparse, subprocess
from mutagen import File as MFile
from mutagen.id3 import ID3, ID3NoHeaderError, TCON
from mutagen.flac import FLAC
from mutagen.mp4 import MP4
from mutagen.oggopus import OggOpus
from mutagen.oggvorbis import OggVorbis


# OGG/OPUS use Vorbis-comment-style tags (lowercase keys, multi-value), same
# shape as FLAC. Treat them as a family so the read/write paths are shared.
VORBIS_LIKE = (FLAC, OggOpus, OggVorbis)

ALEMBIC_CONFIG_DIR = os.environ.get("ALEMBIC_CONFIG_DIR", "/config")
WHITELIST_FILE = f"{ALEMBIC_CONFIG_DIR}/pipeline/genres-whitelist.txt"


MAX_PARTS = 3   # mirrors spotify-genre.py's MAX_GENRES
SEPARATOR = "; "

# Map normalized variant → normalized whitelist entry. Spotify and various
# Soulseek uploaders use spelling variants of the same genre — we want all of
# them to land on a single canonical whitelist string.
ALIASES = {
    "rnb":          "rb",       # Rnb → R&B
    "randb":        "rb",       # R and B
    "hiphop":       "hiphop",   # noop after normalization
    "hiphopalt":    "hiphop",
    "hphop":        "hiphop",
    "drumnbass":    "drumbass",  # Drum n Bass → Drum & Bass
    "drumandbass":  "drumbass",  # Drum and Bass
    "dnb":          "drumbass",  # DnB / D&B
    "kpop":         "kpop",
    "jpop":         "jpop",
    "trance":       "trance",
    "idm":          "idm",
    "edm":          "edm",
    "ukgarage":     "ukgarage",
    "rocknroll":    "rockroll",  # Rock n Roll → Rock & Roll
    "rockandroll":  "rockroll",
    "popfilm":      "popfilm",
    "popsoundtrack": "popfilm",
}


def normed(s):
    """lowercase + strip whitespace/hyphens/underscores/&/slashes/dots."""
    return re.sub(r"[\s\-_&/.]+", "", s).lower()


def load_whitelist():
    """Return dict: normalized form → canonical display string from the file."""
    out = {}
    with open(WHITELIST_FILE) as f:
        for line in f:
            line = line.strip()
            if line and not line.startswith("#"):
                out[normed(line)] = line
    return out


# Splits on:  ;   ,   /   space-dash-space  (any of these = multi-genre delimiter)
SPLIT_RE = re.compile(r"\s+-\s+|[;,/]")


def canonicalize(value, wl):
    """Try to parse value into ≤MAX_PARTS canonical whitelist entries.

    Returns:
      (parts_canonical, all_mapped) where:
        parts_canonical = list of whitelist display strings (deduped, original
                          order), possibly empty
        all_mapped      = True iff every parsed part landed on a whitelist entry
                          AND the part count is within MAX_PARTS
    """
    raw_parts = [p.strip(" \t[]()<>") for p in SPLIT_RE.split(value)]
    raw_parts = [p for p in raw_parts if p]
    if not raw_parts or len(raw_parts) > MAX_PARTS:
        return [], False

    seen = set()
    out = []
    for p in raw_parts:
        n = normed(p)
        # Try alias first, then direct whitelist hit
        canon_norm = ALIASES.get(n, n)
        display = wl.get(canon_norm)
        if not display:
            return [], False
        if display not in seen:
            out.append(display)
            seen.add(display)
    return out, True


def assess(value, wl):
    """Return ('clean'|'rewrite'|'blank', new_value).

    - 'clean'   : current value already canonical, no write needed
    - 'rewrite' : all parts map cleanly; rewrite file as new_value
    - 'blank'   : at least one part is junk; blank the field
    """
    if not value or not value.strip():
        return "clean", None
    canon_parts, ok = canonicalize(value, wl)
    if not ok:
        return "blank", None
    new_value = SEPARATOR.join(canon_parts)
    if new_value == value:
        return "clean", None
    return "rewrite", new_value


def is_genre_locked(m):
    """True if the file has GENRE_LOCK=1 (set by fix-genre.sh).

    Manually-overridden genres often won't match the whitelist (e.g. 'Bass'
    isn't in the whitelist but G Jones genuinely needs it). The lock means
    'this is intentional, leave it alone'."""
    if not m or not m.tags:
        return False
    try:
        if isinstance(m, VORBIS_LIKE):
            val = (m.tags.get("genre_lock") or m.tags.get("GENRE_LOCK") or [None])[0]
            return val == "1"
        if isinstance(m, MP4):
            return False  # fix-genre.sh doesn't set lock on MP4 yet
        for frame in m.tags.getall("TXXX"):
            if frame.desc.upper() == "GENRE_LOCK" and frame.text and frame.text[0] == "1":
                return True
    except (AttributeError, IndexError, TypeError):
        pass
    return False


def get_genre(filepath):
    """Return (mutagen_obj, joined_genre_string).

    Vorbis comments (FLAC) and ID3v2.4 TCON both allow multiple GENRE values
    per file. Beets joins them with '; ' for display. We mirror that so a file
    with 8 separate GENRE tags is correctly seen as an 8-part chain (and gets
    flagged as over-tagged), not just whatever the first tag happened to be.
    """
    try:
        m = MFile(filepath)
    except Exception:
        return None, None
    if not m or not m.tags:
        return m, None
    if isinstance(m, VORBIS_LIKE):
        vals = m.tags.get("genre") or []
        return m, ("; ".join(str(v) for v in vals) if vals else None)
    if isinstance(m, MP4):
        vals = m.tags.get("©gen") or []
        return m, ("; ".join(str(v) for v in vals) if vals else None)
    if "TCON" in m.tags and m.tags["TCON"].text:
        vals = m.tags["TCON"].text
        return m, "; ".join(str(v) for v in vals)
    return m, None


def clear_genre(filepath, m):
    try:
        if isinstance(m, VORBIS_LIKE):
            if "genre" in m.tags:
                del m.tags["genre"]
                m.save()
        elif isinstance(m, MP4):
            if "©gen" in m.tags:
                del m.tags["©gen"]
                m.save()
        else:
            tags = ID3(filepath)
            tags.delall("TCON")
            tags.save(filepath)
        return True
    except (AttributeError, ID3NoHeaderError, Exception):
        return False


def write_genre(filepath, m, value):
    """Overwrite GENRE with a single string value (replacing any multi-value
    set the file might have had — that's the whole point of this rewrite)."""
    try:
        if isinstance(m, VORBIS_LIKE):
            m.tags["genre"] = [value]
            m.save()
        elif isinstance(m, MP4):
            m.tags["©gen"] = [value]
            m.save()
        else:
            try:
                tags = ID3(filepath)
            except ID3NoHeaderError:
                tags = ID3()
            tags.delall("TCON")
            tags.add(TCON(encoding=3, text=value))
            tags.save(filepath)
        return True
    except (AttributeError, ID3NoHeaderError, Exception):
        return False


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--apply", action="store_true")
    args = ap.parse_args()

    wl = load_whitelist()
    print(f"[clear-bad-genres] whitelist has {len(wl)} entries")

    # beets now runs in-process in this same container, so the paths it
    # reports are already valid here directly — no container-to-host
    # translation needed.
    r = subprocess.run(
        ["beet", "ls", "-f", "$path"],
        capture_output=True, text=True, timeout=120
    )
    paths = [p for p in r.stdout.splitlines() if p.strip()]
    print(f"[clear-bad-genres] {len(paths)} tracks to scan")

    rewritten = blanked = failed = locked = 0
    rewrite_examples = {}   # original value -> (sample filename, canonical form)
    blank_examples = {}     # original value -> sample filename

    for p in paths:
        if not os.path.exists(p):
            continue
        m, current = get_genre(p)
        if m is None or current is None:
            continue
        if is_genre_locked(m):
            locked += 1
            continue
        action, new_value = assess(current, wl)
        if action == "clean":
            continue

        if action == "rewrite":
            if current not in rewrite_examples:
                rewrite_examples[current] = (os.path.basename(p), new_value)
            if args.apply:
                if write_genre(p, m, new_value):
                    rewritten += 1
                else:
                    failed += 1
            else:
                rewritten += 1
        else:  # blank
            if current not in blank_examples:
                blank_examples[current] = os.path.basename(p)
            if args.apply:
                if clear_genre(p, m):
                    blanked += 1
                else:
                    failed += 1
            else:
                blanked += 1

    verb = "rewrote" if args.apply else "would-rewrite"
    bverb = "blanked" if args.apply else "would-blank"
    print(f"\n[clear-bad-genres] {verb}: {rewritten} tracks, {bverb}: {blanked} tracks, "
          f"locked-skipped: {locked}, failed: {failed}")

    if rewrite_examples:
        print(f"\nNormalize-rewrites ({len(rewrite_examples)} distinct):")
        for old, (sample, new) in sorted(rewrite_examples.items())[:50]:
            print(f"  {old!r:40s}  →  {new!r:30s}  e.g. {sample}")
        if len(rewrite_examples) > 50:
            print(f"  ...and {len(rewrite_examples)-50} more")

    if blank_examples:
        print(f"\nBlanked (will be refilled by spotify-genre) ({len(blank_examples)} distinct):")
        for value, sample in sorted(blank_examples.items())[:50]:
            print(f"  {value!r:60s}  e.g. {sample}")
        if len(blank_examples) > 50:
            print(f"  ...and {len(blank_examples)-50} more")

    if args.apply and (rewritten or blanked):
        print("\n[clear-bad-genres] now run:")
        print("  beet update")
        if blanked:
            print("  spotify-genre.py --apply --force")


if __name__ == "__main__":
    sys.exit(main())