#!/usr/bin/env python3
"""
bandcamp-tag — Tag Bandcamp WAV downloads using Bandcamp metadata.

For each 'Artist - Album' folder in MUSIC_DIR:
  1. Checks which ID3 tags are missing (year, genre, track total, art)
  2. Searches Bandcamp for the release (or uses --url)
  3. Fetches TralbumData JSON from the Bandcamp album page
  4. Writes missing tags and embeds album art using mutagen

Usage:
  bandcamp-tag.py [MUSIC_DIR]               # scan all albums
  bandcamp-tag.py --album "Artist - Title"  # one album
  bandcamp-tag.py --album "..." --url https://artist.bandcamp.com/album/slug
  bandcamp-tag.py --dry-run                 # show plan without writing
  bandcamp-tag.py --force                   # rewrite tags even if present
  bandcamp-tag.py --scan                    # just list what needs updating

Dependencies: mutagen, requests  (both already installed)
"""

import argparse
import html
import json
import os
import re
import sys
import time
import xml.etree.ElementTree as ET
from pathlib import Path

import requests
from mutagen.id3 import (
    APIC, TALB, TCON, TDRC, TIT2, TPOS, TPE1, TPE2, TRCK,
)
from mutagen.wave import WAVE

# ── Config ────────────────────────────────────────────────────────────────────

DEFAULT_MUSIC_DIR = Path('/mnt/local/music')
REQUEST_DELAY     = 0.8   # seconds between Bandcamp requests
ART_MAX_BYTES     = 3 * 1024 * 1024  # skip art larger than 3 MB

SESSION = requests.Session()
SESSION.headers['User-Agent'] = (
    'Mozilla/5.0 (X11; Linux aarch64) AppleWebKit/537.36 '
    '(KHTML, like Gecko) Chrome/120.0.0.0 Safari/537.36'
)

# ── Bandcamp fetch ────────────────────────────────────────────────────────────

def _bc_slug(s: str) -> str:
    """Convert a string to a Bandcamp-style URL slug."""
    s = s.lower().strip()
    s = re.sub(r'[^a-z0-9 -]', '', s)
    s = re.sub(r'\s+', '-', s)
    return s.strip('-')


def _bc_subdomains(artist: str) -> list[str]:
    """Generate plausible Bandcamp subdomain candidates for an artist name."""
    base = re.sub(r'[^a-z0-9]', '', artist.lower())   # no-spaces, no-hyphens
    hyph = _bc_slug(artist)                             # with hyphens
    candidates = []
    for v in (base, hyph, base + 'music', base + 'dnb',
              hyph + 'music', base.replace('-', '')):
        if v and v not in candidates:
            candidates.append(v)
    return candidates


def bandcamp_search(artist: str, album: str) -> str | None:
    """
    Return the URL of the best Bandcamp album match, or None.

    Strategy: generate plausible subdomain + album-slug combinations and
    probe them with HEAD requests (fast, no JS rendering needed).
    """
    album_slug = _bc_slug(album)
    subdomains  = _bc_subdomains(artist)

    for sub in subdomains:
        url = f'https://{sub}.bandcamp.com/album/{album_slug}'
        try:
            r = SESSION.head(url, timeout=8, allow_redirects=True)
            if r.status_code == 200:
                return str(r.url)
        except Exception:
            pass

    return None


def fetch_tralbum(page_url: str) -> dict | None:
    """Fetch and parse TralbumData JSON from a Bandcamp album page."""
    try:
        r = SESSION.get(page_url, timeout=20)
        r.raise_for_status()
    except Exception as e:
        print(f'  [fetch error] {e}')
        return None

    text = r.text
    data = None

    # Primary: data-tralbum attribute (HTML-entity-encoded JSON)
    m = re.search(r'data-tralbum="([^"]*)"', text)
    if m:
        try:
            data = json.loads(html.unescape(m.group(1)))
        except json.JSONDecodeError:
            pass

    # Fallback: inline TralbumData variable
    if data is None:
        m = re.search(r'TralbumData\s*=\s*(\{.*?\});\s*(?:\n|$)', text, re.DOTALL)
        if m:
            try:
                data = json.loads(m.group(1))
            except json.JSONDecodeError:
                pass

    # Fallback: JSON-LD
    if data is None:
        for m in re.finditer(
            r'<script[^>]+type="application/ld\+json"[^>]*>(.*?)</script>',
            text, re.DOTALL,
        ):
            try:
                ld = json.loads(m.group(1))
                if ld.get('@type') == 'MusicAlbum':
                    data = _ld_to_tralbum(ld)
                    break
            except (json.JSONDecodeError, KeyError):
                pass

    if data is None:
        return None

    # If tags are missing from TralbumData, scrape them from the page HTML
    if not (data.get('tags') or (data.get('current') or {}).get('tags')):
        page_tags = re.findall(r'class="tag"[^>]*>([^<]+)<', text)
        if page_tags:
            data['tags'] = [t.strip() for t in page_tags if t.strip()]

    return data


def _ld_to_tralbum(ld: dict) -> dict:
    """Convert JSON-LD MusicAlbum to a minimal TralbumData-style dict."""
    tracks = []
    for item in (ld.get('track') or {}).get('itemListElement') or []:
        rec = item.get('item', {})
        tracks.append({
            'title': rec.get('name', ''),
            'track_num': item.get('position', len(tracks) + 1),
            'duration': None,
            'disc_number': 1,
        })
    pub = ld.get('datePublished', '')
    return {
        'title': ld.get('name', ''),
        'artist': (ld.get('byArtist') or {}).get('name', ''),
        'art_id': None,
        'current': {
            'publish_date': pub,
            'tags': [],
        },
        'tags': [],
        'trackinfo': tracks,
    }


def fetch_art(art_id: int) -> tuple[bytes | None, str]:
    """Download album art from Bandcamp CDN. Returns (bytes, mime) or (None, '')."""
    for size in ('10', '0'):      # 10 = 1200px max, 0 = original
        url = f'https://f4.bcbits.com/img/a{art_id}_{size}.jpg'
        try:
            r = SESSION.get(url, timeout=30, stream=True)
            if not r.ok:
                continue
            mime = r.headers.get('content-type', 'image/jpeg').split(';')[0].strip()
            if not mime.startswith('image/'):
                continue
            data = r.content
            if len(data) > ART_MAX_BYTES:
                continue
            if len(data) > 1000:
                return data, mime
        except Exception:
            pass
    return None, ''


# ── NFO helper ────────────────────────────────────────────────────────────────

def parse_nfo(nfo_path: Path) -> dict:
    """Parse a Kodi/Jellyfin album.nfo into a minimal TralbumData-style dict."""
    try:
        tree = ET.parse(nfo_path)
        root = tree.getroot()

        def text(tag):
            el = root.find(tag)
            return el.text.strip() if el is not None and el.text else ''

        tracks = []
        for t in root.findall('track'):
            pos_el = t.find('position')
            title_el = t.find('title')
            pos = int(pos_el.text) if pos_el is not None and pos_el.text else len(tracks) + 1
            title = title_el.text.strip() if title_el is not None and title_el.text else ''
            tracks.append({'title': title, 'track_num': pos, 'duration': None, 'disc_number': 1})

        return {
            'title': text('title'),
            'artist': text('albumartist') or text('artist'),
            'art_id': None,
            'current': {'publish_date': text('year'), 'tags': [text('genre')] if text('genre') else []},
            'tags': [text('genre')] if text('genre') else [],
            'trackinfo': tracks,
        }
    except Exception:
        return {}


# ── File helpers ──────────────────────────────────────────────────────────────

def parse_folder_name(name: str) -> tuple[str, str] | tuple[None, None]:
    """Split 'Artist - Album' folder name into (artist, album)."""
    parts = name.split(' - ', 1)
    if len(parts) == 2:
        return parts[0].strip(), parts[1].strip()
    return None, None


def parse_track_filename(stem: str, folder_name: str) -> tuple[int | None, str]:
    """
    Return (track_number, title_guess) from a WAV filename stem.
    Input examples:
      'ALIX PEREZ - Gloom EP - 01 Psychosis'
      '01 Psychosis'
    """
    # Strip leading 'Artist - Album - '
    prefix = folder_name + ' - '
    if stem.startswith(prefix):
        stem = stem[len(prefix):]

    m = re.match(r'^(\d{1,3})\s+(.+)$', stem)
    if m:
        return int(m.group(1)), m.group(2).strip()
    return None, stem.strip()


def tag_completeness(wav_path: Path) -> dict:
    """Return dict of missing tag names for a WAV file."""
    missing = {}
    try:
        f = WAVE(wav_path)
        tags = f.tags or {}

        def has(prefix):
            return any(k.startswith(prefix) for k in tags.keys())

        def val(key):
            v = tags.get(key)
            return str(v[0]) if v else ''

        if not val('TIT2'):
            missing['title'] = True
        if not val('TALB'):
            missing['album'] = True
        if not val('TPE1'):
            missing['artist'] = True
        if not val('TDRC'):
            missing['year'] = True
        if not val('TCON'):
            missing['genre'] = True
        if not has('APIC'):
            missing['art'] = True
        trck = val('TRCK')
        if trck and '/' not in trck:
            missing['track_total'] = True
    except Exception:
        missing['read_error'] = True
    return missing


def load_folder_art(folder: Path) -> tuple[bytes | None, str]:
    """Try to load cover art from a folder."""
    candidates = ['cover.jpg', 'cover.jpeg', 'cover.png', 'folder.jpg', 'folder.png']
    for name in candidates:
        p = folder / name
        if p.exists():
            try:
                data = p.read_bytes()
                mime = 'image/jpeg' if name.endswith(('.jpg', '.jpeg')) else 'image/png'
                return data, mime
            except Exception:
                pass
    return None, ''


def write_tags(
    wav_path: Path,
    bc_track: dict,
    bc_album: dict,
    art_bytes: bytes | None,
    art_mime: str,
    track_total: int,
    missing: dict,
    dry_run: bool,
    force: bool,
) -> None:
    title     = bc_track.get('title', '')
    track_num = bc_track.get('track_num') or 1
    disc_num  = bc_track.get('disc_number') or 1

    current   = bc_album.get('current') or {}
    album_str = current.get('title') or bc_album.get('title', '')
    artist    = bc_track.get('artist') or bc_album.get('artist', '')
    tags_list = current.get('tags') or bc_album.get('tags') or []
    genre     = tags_list[0] if tags_list else ''

    pub       = current.get('publish_date', '')
    year      = re.search(r'\d{4}', pub).group(0) if re.search(r'\d{4}', pub) else ''

    if dry_run:
        updates = []
        if force or missing.get('title'):   updates.append(f'title={title!r}')
        if force or missing.get('album'):   updates.append(f'album={album_str!r}')
        if force or missing.get('artist'):  updates.append(f'artist={artist!r}')
        if force or missing.get('year'):    updates.append(f'year={year!r}')
        if force or missing.get('genre'):   updates.append(f'genre={genre!r}')
        if force or missing.get('track_total'): updates.append(f'trck={track_num}/{track_total}')
        if (force or missing.get('art')) and art_bytes: updates.append('art')
        print(f'    [dry] {wav_path.name}: {", ".join(updates) or "no changes"}')
        return

    try:
        f = WAVE(wav_path)
        if f.tags is None:
            f.add_tags()
        t = f.tags

        def set_tag(key, cls, **kw):
            for k in list(t.keys()):
                if k == key or k.startswith(key + ':'):
                    del t[k]
            t[key] = cls(**kw)

        if force or missing.get('title') and title:
            set_tag('TIT2', TIT2, text=[title])
        if force or missing.get('album') and album_str:
            set_tag('TALB', TALB, text=[album_str])
        if force or missing.get('artist') and artist:
            set_tag('TPE1', TPE1, text=[artist])
            set_tag('TPE2', TPE2, text=[bc_album.get('artist', artist)])
        if (force or missing.get('year')) and year:
            set_tag('TDRC', TDRC, text=[year])
        if (force or missing.get('genre')) and genre:
            set_tag('TCON', TCON, text=[genre])
        if force or missing.get('track_total'):
            set_tag('TRCK', TRCK, text=[f'{track_num}/{track_total}'])
        if disc_num > 1:
            set_tag('TPOS', TPOS, text=[str(disc_num)])
        if (force or missing.get('art')) and art_bytes:
            for k in list(t.keys()):
                if k.startswith('APIC'):
                    del t[k]
            t['APIC:cover'] = APIC(
                mime=art_mime, type=3, desc='cover', data=art_bytes,
            )

        f.save()
        print(f'    OK  {wav_path.name}')
    except Exception as e:
        print(f'    ERR {wav_path.name}: {e}')


# ── Title similarity (for fuzzy track matching) ───────────────────────────────

def _jaccard(a: str, b: str) -> float:
    wa = set(re.findall(r'\w+', a.lower()))
    wb = set(re.findall(r'\w+', b.lower()))
    if not wa or not wb:
        return 0.0
    return len(wa & wb) / len(wa | wb)


# ── Album processing ──────────────────────────────────────────────────────────

def process_album(
    folder: Path,
    url: str | None = None,
    force: bool = False,
    dry_run: bool = False,
    scan_only: bool = False,
) -> bool:
    """Tag all WAV files in one album folder. Returns True if any work was done."""
    wav_files = sorted(folder.glob('*.wav'))
    if not wav_files:
        return False

    folder_name = folder.name
    artist, album = parse_folder_name(folder_name)
    if not artist:
        print(f'SKIP  {folder_name}  (can\'t parse artist/album)')
        return False

    # Check which tags are missing across all WAVs
    all_missing: dict[str, bool] = {}
    for wf in wav_files:
        for k, v in tag_completeness(wf).items():
            if v:
                all_missing[k] = True

    if not force and not all_missing:
        if not scan_only:
            print(f'OK    {folder_name}')
        return False

    missing_str = ', '.join(sorted(all_missing)) or 'none'
    print(f'NEED  {folder_name}  [{missing_str}]')

    if scan_only:
        return True

    # ── Load metadata source ──────────────────────────────────────────────────

    bc_data: dict | None = None

    # Try .nfo first (no network needed for title/genre/tracks)
    nfo_path = folder / 'album.nfo'
    if nfo_path.exists():
        nfo = parse_nfo(nfo_path)
        if nfo.get('trackinfo'):
            print(f'  Using .nfo: {nfo_path.name}')
            bc_data = nfo

    # Fetch from Bandcamp if still missing year, genre, or we have no data yet
    still_need_bc = (
        bc_data is None
        or (force or all_missing.get('year')) and not (bc_data.get('current') or {}).get('publish_date', '').strip()
        or (force or all_missing.get('genre')) and not ((bc_data.get('current') or {}).get('tags') or bc_data.get('tags'))
        or all_missing.get('art') and bc_data.get('art_id') is None
    )

    if still_need_bc:
        if not url:
            print(f'  Searching Bandcamp for {artist!r} — {album!r} …')
            url = bandcamp_search(artist, album)
            time.sleep(REQUEST_DELAY)
            if not url:
                print(f'  No Bandcamp result. Use --url to provide one.')
                if bc_data is None:
                    return False
            else:
                print(f'  Found: {url}')

        if url:
            fetched = fetch_tralbum(url)
            time.sleep(REQUEST_DELAY)
            if fetched:
                # Merge: prefer Bandcamp data but keep .nfo track list if Bandcamp has none
                if bc_data and bc_data.get('trackinfo') and not fetched.get('trackinfo'):
                    fetched['trackinfo'] = bc_data['trackinfo']
                bc_data = fetched
            elif bc_data is None:
                print(f'  Could not fetch album data.')
                return False

    if bc_data is None:
        return False

    # ── Album art ─────────────────────────────────────────────────────────────

    art_bytes: bytes | None = None
    art_mime = 'image/jpeg'

    if force or all_missing.get('art'):
        art_id = bc_data.get('art_id') or (bc_data.get('current') or {}).get('art_id')
        if art_id:
            print(f'  Downloading art (id={art_id}) …')
            art_bytes, art_mime = fetch_art(art_id)
            time.sleep(REQUEST_DELAY)
            if art_bytes:
                print(f'  Art: {len(art_bytes) // 1024} KB ({art_mime})')
            else:
                print(f'  Art download failed, falling back to folder art')

        if not art_bytes:
            art_bytes, art_mime = load_folder_art(folder)
            if art_bytes:
                print(f'  Using folder art ({art_mime}, {len(art_bytes)//1024} KB)')

    # ── Build track map ───────────────────────────────────────────────────────

    bc_tracks_by_num: dict[int, dict] = {}
    for t in bc_data.get('trackinfo') or []:
        n = t.get('track_num')
        if n:
            bc_tracks_by_num[n] = t

    total = len(bc_data.get('trackinfo') or []) or len(wav_files)

    # ── Tag each file ─────────────────────────────────────────────────────────

    for wav_file in wav_files:
        stem = wav_file.stem
        track_num, title_guess = parse_track_filename(stem, folder_name)
        file_missing = tag_completeness(wav_file)

        if not force and not file_missing:
            print(f'    --  {wav_file.name}  (complete)')
            continue

        # Match to Bandcamp track
        if track_num and track_num in bc_tracks_by_num:
            bc_track = bc_tracks_by_num[track_num]
        elif bc_tracks_by_num:
            # Fuzzy match by title
            best = max(
                bc_tracks_by_num.values(),
                key=lambda t: _jaccard(t.get('title', ''), title_guess),
            )
            bc_track = best
        else:
            bc_track = {'title': title_guess, 'track_num': track_num or 1, 'disc_number': 1}

        write_tags(
            wav_file, bc_track, bc_data,
            art_bytes, art_mime, total,
            file_missing, dry_run, force,
        )

    return True


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

def main() -> None:
    ap = argparse.ArgumentParser(
        description='Tag Bandcamp WAV downloads with metadata and album art.',
        formatter_class=argparse.RawDescriptionHelpFormatter,
    )
    ap.add_argument(
        'music_dir', nargs='?', default=str(DEFAULT_MUSIC_DIR),
        help=f'Music directory to scan (default: {DEFAULT_MUSIC_DIR})',
    )
    ap.add_argument('--album', metavar='NAME',
                    help='Process only this folder name (relative to MUSIC_DIR or absolute)')
    ap.add_argument('--url', metavar='URL',
                    help='Bandcamp album URL (use with --album to skip search)')
    ap.add_argument('--force', action='store_true',
                    help='Rewrite all tags even if already present')
    ap.add_argument('--dry-run', action='store_true',
                    help='Show planned changes without writing anything')
    ap.add_argument('--scan', action='store_true',
                    help='List albums that need updating and exit')
    args = ap.parse_args()

    music_dir = Path(args.music_dir)
    if not music_dir.is_dir():
        sys.exit(f'Error: {music_dir} is not a directory')

    if args.album:
        target = Path(args.album)
        if not target.is_absolute():
            target = music_dir / args.album
        if not target.is_dir():
            sys.exit(f'Error: {target} is not a directory')
        process_album(target, url=args.url, force=args.force,
                      dry_run=args.dry_run, scan_only=args.scan)
        return

    folders = sorted(d for d in music_dir.iterdir() if d.is_dir())
    wav_folders = [d for d in folders if any(d.glob('*.wav'))]
    print(f'Scanning {len(wav_folders)} album folders in {music_dir}\n')

    needs_work = 0
    for folder in wav_folders:
        worked = process_album(folder, force=args.force,
                               dry_run=args.dry_run, scan_only=args.scan)
        if worked:
            needs_work += 1
        if not args.scan:
            time.sleep(0.3)

    if args.scan:
        print(f'\n{needs_work} of {len(wav_folders)} albums need updating.')
    else:
        print(f'\nDone. {needs_work} albums processed.')


if __name__ == '__main__':
    main()