#!/usr/bin/env python3
"""Generate a rekordbox-compatible DJ_PLAYLISTS XML from a music folder tree.

Every directory becomes a playlist. A directory holding both tracks and
subdirectories becomes a folder node with a playlist inside it, so the
rekordbox playlist tree ends up mirroring the folder layout exactly.

Usage:
    python rekordbox_xml.py /path/to/music -o rekordbox.xml

Import in rekordbox:
    Preferences > View > Layout > enable "rekordbox xml"
    Preferences > Advanced > Database > rekordbox xml > select the file
    The tree then appears under "rekordbox xml" in the sidebar.
"""
import argparse
import os
import shutil
import subprocess
import sys
import time
from urllib.parse import quote
from xml.sax.saxutils import quoteattr

try:
    import mutagen
    from mutagen.id3 import ID3
except ImportError:
    sys.exit("mutagen is required:  pip install mutagen")

AUDIO_EXT = (".wav", ".mp3", ".flac", ".m4a", ".aiff", ".aif", ".wma",
             ".ogg", ".opus", ".alac")
KIND = {".wav": "WAV File", ".mp3": "MP3 File", ".flac": "FLAC File",
        ".m4a": "M4A File", ".aiff": "AIFF File", ".aif": "AIFF File",
        ".wma": "WMA File", ".ogg": "OGG File", ".opus": "OPUS File",
        ".alac": "ALAC File"}

# ID3 frames (MP3, WAV, AIFF) -> internal field name
ID3MAP = {"TIT2": "title", "TPE1": "artist", "TALB": "album", "TCON": "genre",
          "TDRC": "year", "TRCK": "track", "TBPM": "bpm", "TKEY": "key",
          "TPOS": "disc", "TCOM": "composer", "TPUB": "label",
          "TPE2": "albumartist", "TPE4": "remixer"}
# Vorbis comments (FLAC, OGG)
VORBIS = {"title": "title", "artist": "artist", "album": "album",
          "genre": "genre", "date": "year", "tracknumber": "track",
          "bpm": "bpm", "initialkey": "key", "discnumber": "disc",
          "composer": "composer", "organization": "label",
          "albumartist": "albumartist", "remixer": "remixer"}
# MP4 atoms (M4A)
MP4MAP = {"\xa9nam": "title", "\xa9ART": "artist", "\xa9alb": "album",
          "\xa9gen": "genre", "\xa9day": "year", "\xa9wrt": "composer",
          "aART": "albumartist", "tmpo": "bpm"}

HAVE_FFPROBE = shutil.which("ffprobe") is not None
NOWIN = 0x08000000 if os.name == "nt" else 0


def xmlsafe(s):
    """Drop characters XML 1.0 forbids; control codes leak in from sloppy tags."""
    return "".join(c for c in str(s)
                   if c in "\t\n\r"
                   or 0x20 <= ord(c) <= 0xD7FF
                   or 0xE000 <= ord(c) <= 0xFFFD
                   or 0x10000 <= ord(c) <= 0x10FFFF)


def first(v):
    if isinstance(v, (list, tuple)):
        return str(v[0]) if v else ""
    return str(v) if v is not None else ""


def ffprobe_length(path):
    """Duration fallback for files mutagen cannot measure."""
    if not HAVE_FFPROBE:
        return 0
    try:
        r = subprocess.run(
            ["ffprobe", "-v", "error", "-show_entries", "format=duration",
             "-of", "default=nw=1:nk=1", "--", path],
            capture_output=True, text=True, encoding="utf-8",
            errors="replace", creationflags=NOWIN, timeout=30)
        return int(float(r.stdout.strip()))
    except Exception:
        return 0


def read_meta(path, ext):
    """Return (fields, length_sec, bitrate_kbps, samplerate)."""
    f = {}
    length = bitrate = samplerate = 0
    try:
        m = mutagen.File(path)
    except Exception:
        m = None
    if m is not None:
        info = getattr(m, "info", None)
        if info is not None:
            length = int(getattr(info, "length", 0) or 0)
            bitrate = int((getattr(info, "bitrate", 0) or 0) / 1000)
            samplerate = int(getattr(info, "sample_rate", 0) or 0)
        tags = m.tags
        if tags is not None:
            if isinstance(tags, ID3) or hasattr(tags, "getall"):
                for frame, field in ID3MAP.items():
                    if frame in tags:
                        fr = tags[frame]
                        f[field] = first(fr.text if hasattr(fr, "text") else fr)
            elif ext == ".m4a":
                for k, field in MP4MAP.items():
                    if k in tags:
                        f[field] = first(tags[k])
            else:
                for k, field in VORBIS.items():
                    if k in tags:
                        f[field] = first(tags[k])
    if not length:
        length = ffprobe_length(path)
    if not bitrate and samplerate:
        bitrate = int(samplerate * 16 * 2 / 1000)
    return f, length, bitrate, samplerate


def location(path):
    """rekordbox expects a percent-encoded file://localhost/ URL."""
    p = os.path.abspath(path).replace("\\", "/")
    if not p.startswith("/"):
        p = "/" + p                      # Windows drive letters
    return "file://localhost" + quote(p, safe="/:")


def clean_num(s, maxv=None):
    s = (s or "").split("/")[0].strip()
    try:
        n = int(float(s))
    except ValueError:
        return 0
    if maxv and not (0 < n <= maxv):
        return 0
    return max(n, 0)


def scan(root, keep_empty=False, keep_apple=False):
    """Map every directory to its playable audio files."""
    tree = {}
    for dp, dns, fns in os.walk(root):
        if not keep_apple:
            dns[:] = [d for d in dns if d != "__MACOSX"]
        dns.sort()
        keep = []
        for x in fns:
            if os.path.splitext(x)[1].lower() not in AUDIO_EXT:
                continue
            if not keep_apple and x.startswith("._"):
                continue             # AppleDouble metadata stub, not audio
            if not keep_empty:
                try:
                    if os.path.getsize(os.path.join(dp, x)) == 0:
                        continue     # failed download
                except OSError:
                    continue
            keep.append(x)
        tree[dp] = sorted(keep)
    return tree


def build(root, out_path, tree, quiet=False):
    subdirs = {}
    for dp in tree:
        parent = os.path.dirname(dp)
        if parent in tree and parent != dp:
            subdirs.setdefault(parent, []).append(dp)
    for k in subdirs:
        subdirs[k].sort()

    cache = {}

    def has_audio(dp):
        if dp not in cache:
            cache[dp] = (bool(tree.get(dp))
                         or any(has_audio(s) for s in subdirs.get(dp, [])))
        return cache[dp]

    tracks, ids, tid = [], {}, 0
    total = sum(len(v) for v in tree.values())
    t0 = time.time()
    for dp in sorted(tree):
        for name in tree[dp]:
            p = os.path.join(dp, name)
            ext = os.path.splitext(name)[1].lower()
            try:
                size = os.path.getsize(p)
                mtime = time.strftime("%Y-%m-%d",
                                      time.localtime(os.path.getmtime(p)))
            except OSError:
                continue
            meta, length, bitrate, sr = read_meta(p, ext)
            tid += 1
            ids[p] = tid
            tracks.append({
                "id": tid, "path": p, "size": size, "date": mtime,
                "kind": KIND.get(ext, "Unknown"), "len": length,
                "br": bitrate, "sr": sr,
                "title": meta.get("title") or os.path.splitext(name)[0],
                "artist": meta.get("artist", ""), "album": meta.get("album", ""),
                "genre": meta.get("genre", ""),
                "composer": meta.get("composer", ""),
                "label": meta.get("label", ""),
                "remixer": meta.get("remixer", ""),
                "year": clean_num(meta.get("year", "")[:4], 3000),
                "track": clean_num(meta.get("track", ""), 999),
                "disc": clean_num(meta.get("disc", ""), 99),
                "bpm": meta.get("bpm", ""), "key": meta.get("key", ""),
            })
            if not quiet and len(tracks) % 2500 == 0:
                print(f"  {len(tracks)}/{total} scanned "
                      f"({time.time() - t0:.0f}s)", flush=True)

    lines = ['<?xml version="1.0" encoding="UTF-8"?>',
             '<DJ_PLAYLISTS Version="1.0.0">',
             '  <PRODUCT Name="rekordbox" Version="6.0.0" Company="AlphaTheta"/>',
             f'  <COLLECTION Entries="{len(tracks)}">']
    for t in tracks:
        try:
            b = float(str(t["bpm"]).split("/")[0])
            bpm = f"{b:.2f}" if 0 < b < 400 else ""
        except (ValueError, TypeError):
            bpm = ""
        attrs = [
            ("TrackID", str(t["id"])), ("Name", t["title"]),
            ("Artist", t["artist"]), ("Composer", t["composer"]),
            ("Album", t["album"]), ("Grouping", ""), ("Genre", t["genre"]),
            ("Kind", t["kind"]), ("Size", str(t["size"])),
            ("TotalTime", str(t["len"])), ("DiscNumber", str(t["disc"])),
            ("TrackNumber", str(t["track"])), ("Year", str(t["year"])),
            ("AverageBpm", bpm), ("DateAdded", t["date"]),
            ("BitRate", str(t["br"])), ("SampleRate", str(t["sr"])),
            ("Comments", ""), ("PlayCount", "0"), ("Rating", "0"),
            ("Location", location(t["path"])), ("Remixer", t["remixer"]),
            ("Tonality", t["key"]), ("Label", t["label"]), ("Mix", ""),
        ]
        lines.append("    <TRACK "
                     + " ".join(f"{k}={quoteattr(xmlsafe(v))}" for k, v in attrs)
                     + "/>")
    lines.append("  </COLLECTION>")

    def playlist_node(dp, files, name, depth):
        ind = "    " + "  " * depth
        out = [f"{ind}<NODE Name={quoteattr(xmlsafe(name))} Type=\"1\" "
               f"KeyType=\"0\" Entries=\"{len(files)}\">"]
        for fn in files:
            k = ids.get(os.path.join(dp, fn))
            if k:
                out.append(f'{ind}  <TRACK Key="{k}"/>')
        out.append(f"{ind}</NODE>")
        return out

    def node(dp, depth):
        files = tree.get(dp, [])
        subs = [s for s in subdirs.get(dp, []) if has_audio(s)]
        name = os.path.basename(dp) or "Playlist"
        ind = "    " + "  " * depth
        out = []
        if files and subs:
            out.append(f"{ind}<NODE Name={quoteattr(xmlsafe(name))} Type=\"0\" "
                       f"Count=\"{len(subs) + 1}\">")
            out += playlist_node(dp, files, name, depth + 1)
            for s in subs:
                out += node(s, depth + 1)
            out.append(f"{ind}</NODE>")
        elif subs:
            out.append(f"{ind}<NODE Name={quoteattr(xmlsafe(name))} Type=\"0\" "
                       f"Count=\"{len(subs)}\">")
            for s in subs:
                out += node(s, depth + 1)
            out.append(f"{ind}</NODE>")
        elif files:
            out += playlist_node(dp, files, name, depth)
        return out

    top = [s for s in subdirs.get(root, []) if has_audio(s)]
    root_files = tree.get(root, [])
    lines.append("  <PLAYLISTS>")
    lines.append(f'    <NODE Type="0" Name="ROOT" '
                 f'Count="{len(top) + (1 if root_files else 0)}">')
    if root_files:
        lines += playlist_node(root, root_files,
                               os.path.basename(root) or "Root", 1)
    for d in top:
        lines += node(d, 1)
    lines.append("    </NODE>")
    lines.append("  </PLAYLISTS>")
    lines.append("</DJ_PLAYLISTS>")

    with open(out_path, "w", encoding="utf-8", newline="\r\n") as fh:
        fh.write("\n".join(lines))

    return len(tracks), sum(1 for x in lines if 'Type="1"' in x)


def main():
    ap = argparse.ArgumentParser(
        description="Generate a rekordbox XML from a music folder tree.")
    ap.add_argument("music_dir", help="root folder to scan")
    ap.add_argument("-o", "--output", default="rekordbox.xml",
                    help="output XML path (default: rekordbox.xml)")
    ap.add_argument("--keep-empty", action="store_true",
                    help="include 0-byte files (skipped by default)")
    ap.add_argument("--keep-apple", action="store_true",
                    help="include __MACOSX / ._ AppleDouble stubs")
    ap.add_argument("-q", "--quiet", action="store_true")
    args = ap.parse_args()

    root = os.path.abspath(args.music_dir)
    if not os.path.isdir(root):
        sys.exit(f"not a directory: {root}")
    if not HAVE_FFPROBE and not args.quiet:
        print("note: ffprobe not found - duration fallback disabled")

    t0 = time.time()
    tree = scan(root, args.keep_empty, args.keep_apple)
    n_tracks, n_pl = build(root, args.output, tree, args.quiet)

    print(f"\nTracks    : {n_tracks}")
    print(f"Playlists : {n_pl}")
    print(f"Output    : {args.output} "
          f"({os.path.getsize(args.output) / 2**20:.1f} MB)")
    print(f"Elapsed   : {time.time() - t0:.0f}s")


if __name__ == "__main__":
    main()