#!/usr/bin/env python3
"""
mpc60build - turn a folder tree of WAVs into a finished MPC60 drive image.

    HDA/
      Roland R-70/        -> partition A   (subfolders inside are just grouping)
      AcidLab Miami/      -> partition B
      LSDJ/               -> partition C

    ./mpc60build.py HDA -o drive.hda

Each top-level folder becomes one lettered partition. Its WAVs are converted
to 40 kHz / 12-bit .SND and packed into .PGM programs automatically.

THE LIMIT THAT ACTUALLY BITES is not disk space - a partition holds 29 MB -
but the MPC60's 750 KB of sound memory, about 12.8 seconds of audio. So a
program is capped at 15 sounds (all the donor .PGM exposes) AND at a byte
budget. Anything larger is split into as many programs as it takes, each one
balanced across drum roles so every program is playable on its own rather
than being fifteen consecutive kicks.

Requires: numpy, and three reference files taken from a real MPC60 disk -
a boot sector, a .PGM to use as a donor, and any .SND for its header.
"""
import argparse, os, re, struct, sys
import numpy as np

SECT, SPC = 512, 16
RESERVED, NFATS, SPF, ROOT_ENTS = 1, 2, 11, 224
ROOT_SECT = ROOT_ENTS * 32 // SECT
DATA_START = RESERVED + NFATS * SPF + ROOT_SECT
RATE, SLOTS = 40000, 16   # donor exposes 16 usable slots, not 15
NAME_LEN, FNAME_LEN = 16, 8

ROLES = [("KICK", ("kick", "bd", "_k")), ("SNARE", ("snar", "snare", "sd")),
         ("RIM", ("rim",)), ("CLAP", ("clap", "clp")),
         ("HAT", ("hat", "hh", "chh", "ohh")), ("TOM", ("tom",)),
         ("CYM", ("cym", "crash", "ride", "china")),
         ("PERC", ("perc", "tamb", "agogo", "conga", "cow", "clave", "timb",
                   "block", "whis", "shak", "bongo")),
         ("STAB", ("stab",))]


def role(n):
    l = n.lower()
    for i, (_, keys) in enumerate(ROLES):
        if any(k in l for k in keys):
            return i
    return len(ROLES)


# ---------- audio ----------------------------------------------------------

def read_wav(path):
    with open(path, "rb") as f:
        if f.read(12)[:4] != b"RIFF":
            raise ValueError("not RIFF")
        fmt = None
        while True:
            h = f.read(8)
            if len(h) < 8:
                raise ValueError("no data chunk")
            cid, size = struct.unpack("<4sI", h)
            body = f.read(size)
            if size % 2:
                f.read(1)
            if cid == b"fmt ":
                fmt = body
            elif cid == b"data":
                tag, nch, rate, _, _, bits = struct.unpack("<HHIIHH", fmt[:16])
                if tag == 0xFFFE and len(fmt) >= 40:
                    tag = struct.unpack("<H", fmt[24:26])[0]
                return decode(body, tag, bits, nch), rate


def decode(raw, tag, bits, nch):
    if tag == 1 and bits == 8:
        d = (np.frombuffer(raw, np.uint8).astype(np.float32) - 128.0) / 128.0
    elif tag == 1 and bits == 16:
        d = np.frombuffer(raw, "<i2").astype(np.float32) / 32768.0
    elif tag == 1 and bits == 24:
        n = len(raw) // 3 * 3
        b = np.frombuffer(raw[:n], np.uint8).reshape(-1, 3).astype(np.int32)
        v = b[:, 0] | (b[:, 1] << 8) | (b[:, 2] << 16)
        d = np.where(v & 0x800000, v - 0x1000000, v).astype(np.float32) / 8388608.0
    elif tag == 1 and bits == 32:
        d = np.frombuffer(raw, "<i4").astype(np.float32) / 2147483648.0
    elif tag == 3 and bits == 32:
        d = np.frombuffer(raw, "<f4").astype(np.float32)
    elif tag == 3 and bits == 64:
        d = np.frombuffer(raw, "<f8").astype(np.float32)
    else:
        raise ValueError(f"unsupported WAV: tag {tag}, {bits}-bit")
    if nch > 1:
        d = d[:len(d) // nch * nch].reshape(-1, nch).mean(axis=1)
    return d


def resample(x, src, dst):
    """Band-limited via spectral truncation - no aliasing on downsample."""
    if src == dst or len(x) == 0:
        return x
    n_out = int(round(len(x) * dst / src))
    if n_out < 2:
        return np.zeros(0, np.float32)
    X = np.fft.rfft(x.astype(np.float64))
    Y = np.zeros(n_out // 2 + 1, complex)
    k = min(len(X), len(Y))
    Y[:k] = X[:k]
    return (np.fft.irfft(Y, n_out) * (n_out / len(x))).astype(np.float32)


def to_codes(x, gain=1.0, dither=False):
    """Float audio -> 12-bit codes.

    Rounding to 12 bits leaves an error that, at normal levels, is already
    uncorrelated with the signal - so dither buys nothing there and costs
    about 5 dB of SNR. Down in a decay tail below roughly -60 dBFS the error
    starts tracking the signal instead, which is heard as grain. TPDF dither
    trades that grain for steady noise.
    """
    if len(x) == 0:
        return np.zeros(2, np.int32)
    scaled = x * gain * 2047.0
    if dither:
        rng = np.random.default_rng(60)          # fixed, so builds are reproducible
        scaled = scaled + (rng.random(len(scaled)) - rng.random(len(scaled)))
    q = np.clip(np.round(scaled), -2048, 2047).astype(np.int32)
    if len(q) % 2:
        q = np.append(q, q[-1])
    return q


def pack12(s):
    a = s[0::2].astype(np.uint32) & 0xFFF
    b = s[1::2].astype(np.uint32) & 0xFFF
    o = np.empty((len(a), 3), np.uint8)
    o[:, 0] = (a >> 4) & 0xFF
    o[:, 1] = (a & 0x0F) | ((b & 0x0F) << 4)
    o[:, 2] = (b >> 4) & 0xFF
    return o.tobytes()


# ---------- names ----------------------------------------------------------

def clean(text, limit):
    ok = "ABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789-_#"
    s = "".join(c for c in text.upper().replace(" ", "_") if c in ok)
    out = []
    for c in s:
        if c in "-_" and out and out[-1] == "_":
            continue
        out.append("_" if c in "-_" else c)
    return ("".join(out).strip("_") or "SND")[:limit]


def uniq(base, taken, limit):
    if base not in taken:
        taken.add(base)
        return base
    for i in range(1, 10000):
        suf = str(i)
        c = base[:limit - len(suf)] + suf
        if c not in taken:
            taken.add(c)
            return c
    raise RuntimeError("cannot uniquify " + base)


# ---------- filesystem -----------------------------------------------------

# Program layout, decompiled from a real MPC60 PROGRAM_.pgm:
#   0x0002  name table, 16-byte ASCII names, stride 17, 19 slots
#   0x0882  parallel u32 sample-length table, one per slot
#   0x0ac0  pad records, stride 24; byte 0 is the slot index, 0xFF = unassigned
# The two tables are positionally locked: slot i's length lives at 0x0882+4i.
# Writing names without the lengths leaves the donor's sizes in place, which
# tells the MPC to load samples of entirely the wrong length.
PGM_NAME_OFF, PGM_STRIDE, PGM_NSLOTS, PGM_LEN_OFF = 0x02, 17, 19, 0x0882


def donor_slots(donor):
    return [i for i in range(PGM_NSLOTS)
            if donor[PGM_NAME_OFF + PGM_STRIDE * i] not in (0x00, 0x20)]


PGM_PAD_OFF, PGM_PAD_STRIDE, PGM_NPADS = 0x0ac0, 24, 64


def patch_pgm(donor, entries):
    """entries: [(name, sample_count)] placed into the donor's own used slots."""
    b = bytearray(donor)
    cleared = set()
    for k, i in enumerate(donor_slots(donor)):
        o = PGM_NAME_OFF + PGM_STRIDE * i
        if k < len(entries):
            name, count = entries[k]
            b[o:o + NAME_LEN] = name.ljust(NAME_LEN)[:NAME_LEN].encode("ascii")
            struct.pack_into("<I", b, PGM_LEN_OFF + 4 * i, count)
        else:
            b[o:o + NAME_LEN] = b"\x00" * NAME_LEN
            struct.pack_into("<I", b, PGM_LEN_OFF + 4 * i, 0)
            cleared.add(i)
    # A pad still pointing at a slot we just emptied would reference a nameless,
    # zero-length sample. The donor never does this: unassigned pads hold 0xFF.
    for k in range(PGM_NPADS):
        o = PGM_PAD_OFF + PGM_PAD_STRIDE * k
        if o < len(b) and b[o] in cleared:
            b[o] = 0xFF
    return bytes(b)


def build_partition(files, boot, tot):
    """files: [(name.EXT, bytes)] -> one FAT12 volume."""
    ncl = (tot - DATA_START) // SPC
    img = bytearray(tot * SECT)
    img[0:SECT] = boot
    fat = [0] * (ncl + 2)
    fat[0], fat[1] = 0xFF8, 0xFFF
    root, nxt, placed = bytearray(), 2, 0
    for name, blob in files:
        need = max(1, (len(blob) + SPC * SECT - 1) // (SPC * SECT))
        if nxt + need > ncl + 2 or len(root) + 32 > ROOT_ENTS * 32:
            break
        first = nxt
        for i in range(need):
            fat[first + i] = 0xFFF if i == need - 1 else first + i + 1
        o = (DATA_START + (first - 2) * SPC) * SECT
        img[o:o + len(blob)] = blob
        nxt += need
        base, _, ext = name.partition(".")
        e = bytearray(32)
        e[0:11] = (base[:8].ljust(8) + ext[:3].ljust(3)).upper().encode("ascii")
        e[11] = 0x00                      # the MPC60 writes 0, not 0x20
        e[26:28] = struct.pack("<H", first)
        e[28:32] = struct.pack("<I", len(blob))
        root.extend(e)
        placed += 1
    packed = bytearray()
    for i in range(0, len(fat) - len(fat) % 2, 2):
        a, b = fat[i], fat[i + 1]
        packed += bytes([a & 0xFF, ((a >> 8) & 0x0F) | ((b & 0x0F) << 4), (b >> 4) & 0xFF])
    if len(fat) % 2:
        packed += bytes([fat[-1] & 0xFF, (fat[-1] >> 8) & 0x0F])
    packed = bytes(packed).ljust(SPF * SECT, b"\0")[:SPF * SECT]
    for n in range(NFATS):
        o = (RESERVED + n * SPF) * SECT
        img[o:o + SPF * SECT] = packed
    o = (RESERVED + NFATS * SPF) * SECT
    img[o:o + len(root)] = root
    return bytes(img), placed


# ---------- packing --------------------------------------------------------

def pack_programs(sounds, budget):
    """Round-robin across roles so each program is a playable kit."""
    buckets = {}
    for s in sounds:
        buckets.setdefault(role(s["long"]), []).append(s)
    for v in buckets.values():
        v.sort(key=lambda s: len(s["blob"]))
    progs, cur, size = [], [], 0
    order = sorted(buckets)
    while any(buckets[r] for r in order):
        for r in order:
            if not buckets[r]:
                continue
            s = buckets[r].pop(0)
            b = len(s["blob"]) - 39
            if len(cur) >= SLOTS or (cur and size + b > budget):
                progs.append(cur); cur, size = [], 0
            cur.append(s); size += b
    if cur:
        progs.append(cur)
    return progs


def main():
    p = argparse.ArgumentParser(description="WAV folder tree -> MPC60 drive image.")
    p.add_argument("folder", help="the HDA directory; each subfolder = one partition")
    p.add_argument("-o", "--out", default="HD0_512.hda",
                   help="ZuluSCSI only sees files named HD<id>_<blocksize>.hda "
                        "- HD0_512.hda through HD6_512.hda, never HD7")
    p.add_argument("--memory-kb", type=int, default=700,
                   help="per-program budget; 700 for a stock MPC60, 1400 if expanded")
    p.add_argument("--partitions", type=int, default=26)
    p.add_argument("--reserve", type=int, default=0,
                   help="leave the first N partitions empty and formatted, "
                        "for your own sounds and programs")
    p.add_argument("--max-sounds", type=int, default=190,
                   help="per partition; 224 root entries minus room for programs")
    p.add_argument("--boot", required=True, help="512-byte MPC60 boot sector")
    p.add_argument("--pgm-donor", required=True, help="a real MPC60 .PGM")
    p.add_argument("--snd-header", required=True, help="any real .SND, for its 39-byte header")
    p.add_argument("--gain", choices=["kit", "sample", "none"], default="kit",
                   help="kit: one gain per folder, set by its loudest sample - "
                        "keeps the balance between sounds AND uses the 12-bit "
                        "range. sample: peak every sound individually (flattens "
                        "dynamics). none: leave source levels alone.")
    p.add_argument("--dither", action="store_true",
                   help="TPDF dither before the 12-bit round. Cleans up decay "
                        "tails, costs ~5 dB SNR. Try it both ways.")
    a = p.parse_args()

    boot = open(a.boot, "rb").read(SECT)
    tot = struct.unpack("<H", boot[19:21])[0]
    donor = open(a.pgm_donor, "rb").read()
    hdr0 = open(a.snd_header, "rb").read(39)
    budget = a.memory_kb * 1024

    kits = sorted(d for d in os.listdir(a.folder)
                  if os.path.isdir(os.path.join(a.folder, d)) and not d.startswith("."))
    if not kits:
        sys.exit(f"no subfolders in {a.folder} - each subfolder becomes one partition")
    if len(kits) > a.partitions:
        print(f"  note: {len(kits)} folders but only {a.partitions} partitions; "
              f"using the first {a.partitions}")
        kits = kits[:a.partitions]

    # ---- pass 1: convert every kit, so sizes are known before allocating ----
    converted = []
    for kit in kits:
        root = os.path.join(a.folder, kit)
        wavs = sorted(os.path.join(dp, f) for dp, _, fs in os.walk(root)
                      for f in fs if f.lower().endswith(".wav") and not f.startswith("."))
        taken, sounds, audio = set(), [], []
        for w in wavs:
            try:
                data, sr = read_wav(w)
            except Exception as e:
                print(f"    skip {os.path.basename(w)}: {e}", file=sys.stderr)
                continue
            audio.append((w, resample(data, sr, RATE)))
        # one gain for the whole kit: the loudest sample defines it, so the
        # quiet ones stay quiet relative to it instead of being pushed up.
        peaks = [float(np.max(np.abs(y))) if len(y) else 0.0 for _, y in audio]
        kit_peak = max(peaks) if peaks else 0.0
        if a.gain == "kit":
            g = 1.0 / kit_peak if kit_peak > 0 else 1.0
        elif a.gain == "none":
            g = 1.0 / kit_peak if kit_peak > 1 else 1.0
        else:
            g = None                       # per-sample, computed below
        for w, y in audio:
            gg = g if g is not None else (1.0 / max(float(np.max(np.abs(y))), 1e-9))
            codes = to_codes(y, gg, a.dither)
            stem = os.path.splitext(os.path.basename(w))[0]
            fn = uniq(clean(stem, FNAME_LEN), taken, FNAME_LEN)
            hdr = bytearray(hdr0)
            hdr[2:18] = fn.ljust(NAME_LEN)[:NAME_LEN].encode("ascii")
            n = len(codes)
            hdr[19], hdr[20], hdr[21] = n & 0xFF, (n >> 8) & 0xFF, (n >> 16) & 0xFF
            for off in (22, 25, 28):
                v = hdr[off] | (hdr[off+1] << 8) | (hdr[off+2] << 16)
                if v > n:
                    hdr[off], hdr[off+1], hdr[off+2] = n & 0xFF, (n >> 8) & 0xFF, (n >> 16) & 0xFF
            sounds.append({"file": fn + ".SND", "long": fn, "count": n,
                           "blob": bytes(hdr) + pack12(codes)})
        if sounds:
            converted.append((kit, sounds))

    # ---- allocate the drive: one partition each, remainder by size ----
    sizes = [sum(len(s["blob"]) - 39 for _, ss in [(k, ss)] for s in ss)
             for k, ss in converted]
    alloc = [1] * len(converted)
    usable = a.partitions - a.reserve
    spare = usable - len(converted)
    if spare < 0:
        sys.exit(f"{len(converted)} kits but only {usable} partitions "
                 f"after reserving {a.reserve}")
    want = [max(0, -(-sz // budget) - 1) for sz in sizes]      # extra beyond the first
    while spare > 0 and any(want):
        # give the next partition to whoever is most starved, proportionally
        i = max(range(len(converted)),
                key=lambda j: (sizes[j] / (alloc[j] * budget)) if want[j] else 0)
        if not want[i]:
            break
        alloc[i] += 1; want[i] -= 1; spare -= 1
    print(f"  {len(converted)} kits into {usable} partitions "
          f"(reserved {a.reserve} empty, {sum(alloc)} allocated)\n")

    # ---- pass 2: trim each kit to its allocation, then build ----
    blank_part, _ = build_partition([], boot, tot)
    parts = [blank_part] * a.reserve
    for r in range(a.reserve):
        print(f"  {chr(65+r)}  (reserved - empty, formatted, yours)")
    for (kit, sounds), n_parts in zip(converted, alloc):
        cap = n_parts * budget
        total = sum(len(s["blob"]) - 39 for s in sounds)
        dropped = 0
        if total > cap or len(sounds) > a.max_sounds * n_parts:
            buckets = {}
            for s in sounds:
                buckets.setdefault(role(s["long"]), []).append(s)
            for v in buckets.values():
                v.sort(key=lambda s: len(s["blob"]))     # short first: more survive
            keep, size, order = [], 0, sorted(buckets)
            while any(buckets[r] for r in order):
                for r in order:
                    if not buckets[r]:
                        continue
                    s = buckets[r].pop(0); b = len(s["blob"]) - 39
                    if size + b > cap or len(keep) >= a.max_sounds * n_parts:
                        continue
                    keep.append(s); size += b
            dropped = len(sounds) - len(keep)
            sounds = keep

        # Chunk into at most n_parts partitions. Packing waste means a chunk
        # can close early, so without this cap a kit silently takes more than
        # its allocation and starves the kits after it.
        chunks, cur, size, lost = [], [], 0, 0
        for s in sorted(sounds, key=lambda s: role(s["long"])):
            b = len(s["blob"]) - 39
            if cur and size + b > budget:
                chunks.append(cur); cur, size = [], 0
                if len(chunks) >= n_parts:
                    lost = 1
                    break
            cur.append(s); size += b
        if cur and len(chunks) < n_parts:
            chunks.append(cur)
        if lost:
            dropped += sum(1 for _ in sounds) - sum(len(c) for c in chunks)

        for ci, chunk in enumerate(chunks):
            if len(parts) >= a.partitions:
                print(f"      out of partitions; {len(chunk)} sounds dropped")
                break
            progs = pack_programs(chunk, budget)
            files = [(s["file"], s["blob"]) for s in chunk]
            tag = clean(kit, 6)
            for i, pr in enumerate(progs):
                pn = tag if len(progs) == 1 else f"{tag}{i+1}"
                files.append((pn[:8] + ".PGM",
                              patch_pgm(donor, [(x["long"], x["count"]) for x in pr])))
            img, placed = build_partition(files, boot, tot)
            parts.append(img)
            kb = sum(len(s["blob"]) - 39 for s in chunk) / 1024
            label = kit[:22] + (f" [{ci+1}/{len(chunks)}]" if len(chunks) > 1 else "")
            print(f"  {chr(65+len(parts)-1)}  {label:<28} {len(chunk):3} sounds  "
                  f"{len(progs):2} pgm  {kb:6.0f} KB"
                  + (f"   (kit trimmed by {dropped})" if dropped and ci == 0 else ""))

    blank, _ = build_partition([], boot, tot)
    with open(a.out, "wb") as f:
        for i in range(a.partitions):
            f.write(parts[i] if i < len(parts) else blank)
    print(f"\n  {a.out}: {os.path.getsize(a.out)} bytes, {a.partitions} partitions")
    print(f"  per-program budget was {a.memory_kb} KB "
          f"({a.memory_kb*1024/1.5/RATE:.1f}s of audio)")


if __name__ == "__main__":
    main()