#!/usr/bin/env python3
"""
Ardour stem salvage — export per-take per-orbit stems from the single
"Tidal Multi" session into the Montreuil26 standard raw-stem layout:

    Prod/_stems/{date}_{take}/orbit-NN.flac   (48k stereo, lossless)
    Prod/_stems/{date}_{take}/_manifest.json  (take, date, orbits, sources)

Each Take = one performance. Sources live in the Ardour interchange dir,
named  TakeN_<chan>-<region>%<L|R>.wav  (chan = "Tidal 1".."Tidal 12"|Keys|Mic).
Mostly 1 region/orbit (direct L/R join); multi-region channels are concatenated
in region order first. SAFE: lossless, never touches the source session.

Usage:
    ardour_stem_export.py --list                 # list takes (no audio)
    ardour_stem_export.py --take Take3 --dry      # show planned export
    ardour_stem_export.py --take Take3            # export one take
    ardour_stem_export.py --all                   # export all takes
"""
import xml.etree.ElementTree as ET
import argparse, collections, datetime, json, re, subprocess, sys
from pathlib import Path

ARDOUR  = Path("/home/pln/Work/Sound/Ardour/Tidal Multi")
SESSION = ARDOUR / "Tidal Multi.ardour"
AUDIO   = ARDOUR / "interchange" / "Tidal Multi" / "audiofiles"
OUT     = Path("/mnt/freebox/PLN/Work/Sound/Prod/_stems")

NAME = re.compile(r'^(?P<take>.+?)_(?P<ch>Tidal \d+|Keys|Mic)-(?P<reg>\d+)(?:%(?P<lr>[LR]))?\.wav$')

def chan_key(ch):
    m = re.match(r'Tidal (\d+)', ch)
    return f"orbit-{int(m.group(1)):02d}" if m else ch.lower()

def parse():
    root = ET.parse(SESSION).getroot()
    sr = int(root.get('sample-rate') or 48000)
    names = [s.get('name') for s in root.iter('Source') if s.get('name')]
    # take -> chan -> lr('L'/'R'/'M') -> {reg:int -> filename}
    takes = collections.defaultdict(lambda: collections.defaultdict(lambda: collections.defaultdict(dict)))
    for nm in names:
        m = NAME.match(nm)
        if not m:
            continue
        lr = m.group('lr') or 'M'
        takes[m.group('take')][m.group('ch')][lr][int(m.group('reg'))] = nm
    return sr, takes

def ordered(d):              # {reg:name} -> [name,...] by region
    return [d[k] for k in sorted(d)]

def take_date(takes, take):
    for ch in takes[take].values():
        for lr in ch.values():
            for nm in lr.values():
                p = AUDIO / nm
                if p.exists():
                    return datetime.date.fromtimestamp(p.stat().st_mtime).isoformat()
    return "0000-00-00"

def concat_to(files, dst):   # concat mono wavs (region order) -> one wav
    if len(files) == 1:
        return files[0]
    inputs = []
    for f in files: inputs += ["-i", str(AUDIO / f)]
    fc = "".join(f"[{i}:a]" for i in range(len(files))) + f"concat=n={len(files)}:v=0:a=1[o]"
    subprocess.run(["ffmpeg","-y","-loglevel","error",*inputs,"-filter_complex",fc,"-map","[o]",str(dst)],check=True)
    return dst

def export_take(sr, takes, take, dry=False):
    date = take_date(takes, take)
    outdir = OUT / f"{date}_{take}"
    kept = []
    for ch in sorted(takes[take], key=chan_key):
        slots = takes[take][ch]
        Ls, Rs, Ms = ordered(slots.get('L',{})), ordered(slots.get('R',{})), ordered(slots.get('M',{}))
        present = [f for f in (Ls+Rs+Ms) if (AUDIO/f).exists()]
        if not present:
            continue
        out = outdir / f"{chan_key(ch)}.flac"
        kept.append({"channel": ch, "out": out.name, "regions": max(len(Ls),len(Rs),len(Ms)),
                     "stereo": bool(Ls and Rs), "sources": present})
        if dry:
            continue
        outdir.mkdir(parents=True, exist_ok=True)
        tmp = []
        if Ls and Rs:                       # stereo: join concatenated L + R
            l = concat_to(Ls, outdir/"_L.wav"); r = concat_to(Rs, outdir/"_R.wav")
            tmp += [p for p in (l,r) if isinstance(p,Path)]
            subprocess.run(["ffmpeg","-y","-loglevel","error","-i",str(AUDIO/l if isinstance(l,str) else l),
                            "-i",str(AUDIO/r if isinstance(r,str) else r),
                            "-filter_complex","[0:a][1:a]join=inputs=2:channel_layout=stereo[a]",
                            "-map","[a]","-c:a","flac",str(out)],check=True)
        else:                                # mono channel (Keys/Mic or single-side)
            src = concat_to(Ms or Ls or Rs, outdir/"_M.wav")
            subprocess.run(["ffmpeg","-y","-loglevel","error","-i",str(AUDIO/src if isinstance(src,str) else src),
                            "-c:a","flac",str(out)],check=True)
        for t in tmp:
            try: Path(t).unlink()
            except OSError: pass
    if not dry and kept:
        (outdir/"_manifest.json").write_text(json.dumps(
            {"take": take, "date": date, "sample_rate": sr,
             "source": "Ardour 'Tidal Multi' interchange",
             "orbits": [k for k in kept]}, indent=2, ensure_ascii=False))
    return date, outdir, kept

if __name__ == "__main__":
    import xml.etree.ElementTree as ET  # noqa (kept local for clarity)
    ap = argparse.ArgumentParser()
    ap.add_argument("--take"); ap.add_argument("--all", action="store_true")
    ap.add_argument("--list", action="store_true"); ap.add_argument("--dry", action="store_true")
    a = ap.parse_args()
    sr, takes = parse()
    if a.list:
        for tk in sorted(takes):
            print(f"{take_date(takes,tk)}  {tk:10}  channels={len(takes[tk])}")
        sys.exit(0)
    targets = sorted(takes) if a.all else ([a.take] if a.take else [])
    if not targets:
        print("specify --take NAME | --all | --list"); sys.exit(1)
    for tk in targets:
        date, outdir, kept = export_take(sr, takes, tk, dry=a.dry)
        tag = "DRY" if a.dry else "OK "
        print(f"[{tag}] {tk} ({date}) -> {outdir}  : {len(kept)} stems"
              + (f"  e.g. {kept[0]['out']} <- {len(kept[0]['sources'])} src" if kept else ""))
