#!/usr/bin/env python3
"""
mpc60conv - convert WAV files to Akai MPC60 .SND, and patch .PGM programs.

FORMAT NOTES
------------
The .SND layout below was recovered from a real MPC60 file (SWELL1-1.SND)
rather than guessed. What is *proven*:

  * header is 39 bytes; audio follows immediately
  * name at offset 2, 16 bytes, ASCII, SPACE padded
  * sample count at offset 19, u24 little-endian
  * audio is 12-bit LINEAR, packed 2 samples per 3 bytes:
        A = (b0 << 4) | (b1 & 0x0F)
        B = (b1 >> 4) | (b2 << 4)
    Verified by decoding the reference file and re-encoding it to
    byte-identical output over all 192000 body bytes, and by three
    independent agreements: the u24 field reads 128000, (192039-39)/3*2
    is 128000, and 128000/40000 is exactly 3.2000 s.

What is NOT known: the u24 fields at 22, 25 and 28 (1280 / 3199 / 5120 in
the reference file) and the bytes at 31 and 38 (both 100). Those are
copied from a donor header. If the MPC rejects a file, they are the first
place to look. Use `--snd-template` to donate them from your own file.

There is no sample-rate field. The MPC60 is fixed at 40 kHz.

Requires: numpy
"""

import argparse
import os
import struct
import sys

import numpy as np

MPC60_RATE = 40000
HEADER_LEN = 39
NAME_OFF, NAME_LEN = 2, 16
COUNT_OFF = 19          # u24 LE, in samples
UNKNOWN_U24 = (22, 25, 28)   # meaning unknown; may be trim/loop pointers
FILENAME_LEN = 8        # MPC disks are DOS/FAT: 8.3

# 39-byte header lifted from a known-good MPC60 file.
HEADER_TEMPLATE = bytes([
    0x01, 0x01, 0x53, 0x57, 0x45, 0x4c, 0x4c, 0x31, 0x2d, 0x31, 0x20, 0x20, 0x20,
    0x20, 0x20, 0x20, 0x20, 0x20, 0x00, 0x00, 0xf4, 0x01, 0x00, 0x05, 0x00, 0x7f,
    0x0c, 0x00, 0x00, 0x14, 0x00, 0x64, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x64,
])

PGM_NAME_TABLE_OFF = 0x24
PGM_NAME_STRIDE = 17
PGM_HEAD_NAME_OFF = 2


# --- WAV reading -------------------------------------------------------

def read_wav(path):
    """Parse RIFF directly. Handles PCM 8/16/24/32 and IEEE float 32/64."""
    with open(path, "rb") as f:
        riff = f.read(12)
        if riff[:4] != b"RIFF" or riff[8:12] != b"WAVE":
            raise ValueError("not a RIFF/WAVE file")
        fmt = None
        while True:
            hdr = f.read(8)
            if len(hdr) < 8:
                raise ValueError("no data chunk")
            cid, size = struct.unpack("<4sI", hdr)
            body = f.read(size)
            if size % 2:
                f.read(1)
            if cid == b"fmt ":
                fmt = body
            elif cid == b"data":
                if fmt is None:
                    raise ValueError("data chunk before fmt")
                tag, nch, rate, _, _, bits = struct.unpack("<HHIIHH", fmt[:16])
                if tag == 0xFFFE and len(fmt) >= 40:      # EXTENSIBLE
                    tag = struct.unpack("<H", fmt[24:26])[0]
                return _decode_pcm(body, tag, bits, nch), rate


def _decode_pcm(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: format tag {tag}, {bits}-bit")
    if nch > 1:
        d = d[: len(d) // nch * nch].reshape(-1, nch).mean(axis=1)
    return d


def write_wav(path, samples12, rate):
    """Write 12-bit codes out as a 16-bit WAV (scaled up 16x)."""
    pcm = (np.clip(samples12, -2048, 2047) * 16).astype("<i2").tobytes()
    with open(path, "wb") as f:
        f.write(b"RIFF" + struct.pack("<I", 36 + len(pcm)) + b"WAVEfmt ")
        f.write(struct.pack("<IHHIIHH", 16, 1, 1, rate, rate * 2, 2, 16))
        f.write(b"data" + struct.pack("<I", len(pcm)) + pcm)


# --- DSP ---------------------------------------------------------------

def resample(x, src, dst):
    """Band-limited resample via spectral truncation - no aliasing."""
    if src == dst or len(x) == 0:
        return x
    n_in = len(x)
    n_out = int(round(n_in * dst / src))
    if n_out < 1:
        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 / n_in)).astype(np.float32)


def quantize12(x, normalize=True):
    """Return int32 array of 12-bit linear codes, plus the gain applied."""
    if len(x) == 0:
        return np.zeros(0, np.int32), 1.0
    peak = float(np.max(np.abs(x)))
    if peak == 0:
        return np.zeros(len(x), np.int32), 1.0
    gain = 1.0 / peak if normalize else (1.0 / peak if peak > 1.0 else 1.0)
    q = np.round(x * gain * 2047.0)
    return np.clip(q, -2048, 2047).astype(np.int32), gain


def pack12(s):
    """12-bit codes -> packed bytes, 2 samples per 3 bytes."""
    if len(s) % 2:
        s = np.append(s, 0)
    a = s[0::2].astype(np.uint32) & 0xFFF
    b = s[1::2].astype(np.uint32) & 0xFFF
    out = np.empty((len(a), 3), np.uint8)
    out[:, 0] = (a >> 4) & 0xFF
    out[:, 1] = (a & 0x0F) | ((b & 0x0F) << 4)
    out[:, 2] = (b >> 4) & 0xFF
    return out.tobytes()


def unpack12(body):
    """Packed bytes -> 12-bit signed codes."""
    n = len(body) // 3 * 3
    b = np.frombuffer(body[:n], np.uint8).astype(np.int32).reshape(-1, 3)
    a = (b[:, 0] << 4) | (b[:, 1] & 0x0F)
    c = (b[:, 1] >> 4) | (b[:, 2] << 4)
    a = np.where(a & 0x800, a - 0x1000, a)
    c = np.where(c & 0x800, c - 0x1000, c)
    out = np.empty(len(a) * 2, np.int32)
    out[0::2], out[1::2] = a, c
    return out


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

def clean(text, limit):
    """Uppercase, strip illegal chars, collapse separator runs, truncate.

    Source names like "Kick - 1.wav" would otherwise spend half of the
    8-character filename budget on " - ".
    """
    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)
    s = "".join(out).strip("_")
    return (s or "SND")[:limit]


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


# --- SND ---------------------------------------------------------------

def u24(buf, off):
    return buf[off] | (buf[off + 1] << 8) | (buf[off + 2] << 16)


def put_u24(buf, off, val):
    buf[off] = val & 0xFF
    buf[off + 1] = (val >> 8) & 0xFF
    buf[off + 2] = (val >> 16) & 0xFF


def build_header(name, count, template=None, clamp=True):
    hdr = bytearray(template if template else HEADER_TEMPLATE)
    if len(hdr) < HEADER_LEN:
        raise ValueError("template shorter than 39 bytes")
    hdr = hdr[:HEADER_LEN]
    hdr[NAME_OFF:NAME_OFF + NAME_LEN] = name.ljust(NAME_LEN)[:NAME_LEN].encode("ascii")
    put_u24(hdr, COUNT_OFF, count)
    if clamp:
        for off in UNKNOWN_U24:
            if u24(hdr, off) > count:
                put_u24(hdr, off, count)
    return bytes(hdr)


def write_snd(path, name, codes, template=None, clamp=True):
    # The reference file's count is even and body == count//2*3 exactly.
    # Keep that invariant: pad an odd tail by repeating the last sample
    # (25 us at 40 kHz) rather than emitting a half-used final byte triple.
    if len(codes) % 2:
        codes = np.append(codes, codes[-1])
    body = pack12(codes)
    hdr = build_header(name, len(codes), template, clamp)
    with open(path, "wb") as f:
        f.write(hdr)
        f.write(body)
    return len(hdr) + len(body)


def read_snd(path):
    raw = open(path, "rb").read()
    name = raw[NAME_OFF:NAME_OFF + NAME_LEN].decode("ascii", "replace").rstrip()
    count = u24(raw, COUNT_OFF)
    codes = unpack12(raw[HEADER_LEN:])[:count]
    return name, count, codes, raw


# --- PGM ---------------------------------------------------------------

def pgm_slots(buf):
    """Offsets of occupied 16-char name slots in the sample name table."""
    out, empties = [], 0
    off = PGM_NAME_TABLE_OFF
    while off + NAME_LEN <= len(buf) and empties < 8:
        field = buf[off:off + NAME_LEN]
        if field[0] not in (0x00, 0x20):
            out.append(off)
            empties = 0
        else:
            empties += 1   # donors have internal gaps; only a long run ends it
        off += PGM_NAME_STRIDE
    return out


def patch_pgm(template, names):
    """Substitute names into the donor program's occupied slots."""
    buf = bytearray(template)
    slots = pgm_slots(buf)
    n = min(len(slots), len(names))
    for i in range(n):
        buf[slots[i]:slots[i] + NAME_LEN] = names[i].ljust(NAME_LEN)[:NAME_LEN].encode("ascii")
    if names:
        buf[PGM_HEAD_NAME_OFF:PGM_HEAD_NAME_OFF + NAME_LEN] = \
            names[0].ljust(NAME_LEN)[:NAME_LEN].encode("ascii")
    return bytes(buf), len(slots), n


# --- commands ----------------------------------------------------------

def collect_wavs(root):
    hits = []
    for dirpath, _, files in os.walk(root):
        for fn in sorted(files):
            if fn.lower().endswith(".wav") and not fn.startswith("."):
                hits.append(os.path.join(dirpath, fn))
    return sorted(hits)


def cmd_convert(args):
    os.makedirs(args.out, exist_ok=True)
    template = open(args.snd_template, "rb").read() if args.snd_template else None
    if template:
        print(f"SND header donor: {args.snd_template}")

    wavs = collect_wavs(args.folder)
    if not wavs:
        print(f"no WAV files under {args.folder}", file=sys.stderr)
        return 1

    files, longnames, skipped, total = set(), [], [], 0
    for path in wavs:
        try:
            data, rate = read_wav(path)
        except Exception as e:
            skipped.append((path, str(e)))
            continue

        stem = os.path.splitext(os.path.basename(path))[0]
        parent = os.path.basename(os.path.dirname(path))
        tag = f"{parent}_{stem}" if os.path.dirname(path) != args.folder.rstrip("/") else stem
        if args.strip and tag.upper().startswith(args.strip.upper()):
            tag = tag[len(args.strip):] or stem
        if args.prefix:
            tag = args.prefix + tag

        data = resample(data, rate, args.rate)
        codes, gain = quantize12(data, normalize=not args.no_normalize)
        fname = uniquify(clean(tag, FILENAME_LEN), files, FILENAME_LEN)
        longname = clean(tag, NAME_LEN)

        n = write_snd(os.path.join(args.out, fname + ".SND"),
                      longname, codes, template, not args.no_clamp)
        total += n
        longnames.append(longname)
        db = 20 * np.log10(gain) if gain > 0 else 0
        print(f"  {os.path.relpath(path, args.folder):<28} -> {fname}.SND  "
              f"{len(codes)/args.rate:6.3f}s  {n/1024:6.1f}K  gain {db:+5.1f} dB")

    for p, e in skipped:
        print(f"  SKIP {os.path.relpath(p, args.folder)}: {e}", file=sys.stderr)
    if not longnames:
        return 1

    if args.pgm_template:
        tpl = open(args.pgm_template, "rb").read()
        data, slots, used = patch_pgm(tpl, longnames)
        out = os.path.join(args.out, clean(args.name, FILENAME_LEN) + ".PGM")
        open(out, "wb").write(data)
        print(f"\n  {out}: {used} of {slots} donor slots renamed"
              f"{'' if used >= len(longnames) else f' ({len(longnames)-used} sounds unmapped)'}")
    else:
        print("\n  no --pgm-template: no program written (the PGM pad records "
              "are not decoded well enough to synthesise one)")

    print(f"\n  {len(longnames)} sounds, {total/1024/1024:.2f} MB"
          + (f", {len(skipped)} skipped" if skipped else ""))
    return 0


def cmd_decode(args):
    name, count, codes, raw = read_snd(args.file)
    out = args.out or os.path.splitext(args.file)[0] + ".wav"
    write_wav(out, codes, args.rate)
    print(f"{args.file}: '{name}' {count} samples "
          f"{count/args.rate:.4f}s -> {out}")
    return 0


def cmd_inspect(args):
    raw = open(args.file, "rb").read()
    size = len(raw)
    print(f"{args.file}  ({size} bytes)\n")
    for row in range(0, min(size, 48), 16):
        chunk = raw[row:row + 16]
        print(f"  {row:04x}  {' '.join(f'{b:02x}' for b in chunk):<48}  "
              + "".join(chr(b) if 32 <= b < 127 else "." for b in chunk))

    count = u24(raw, COUNT_OFF)
    body = size - HEADER_LEN
    print(f"\n  name        @2   '{raw[2:18].decode('ascii', 'replace')}'")
    print(f"  count       @19  {count}")
    print(f"  body        {body} bytes -> {body // 3 * 2} samples")
    ok = body % 3 == 0 and body // 3 * 2 == count
    print(f"  consistent: {'YES' if ok else 'NO - header/body disagree'}")
    print(f"  duration    {count / MPC60_RATE:.4f} s at 40 kHz")
    print("\n  unknown fields (copied verbatim by this tool):")
    for off in UNKNOWN_U24:
        print(f"    u24 @{off}  = {u24(raw, off)}")
    print(f"    u8  @31 = {raw[31]}    u8 @38 = {raw[38]}")
    return 0


def cmd_verify(args):
    """Round-trip a real SND through decode+encode and prove byte equality."""
    name, count, codes, raw = read_snd(args.file)
    body = raw[HEADER_LEN:]
    re_body = pack12(codes)
    exact = re_body == body[:len(re_body)]
    print(f"{args.file}")
    print(f"  name '{name}', {count} samples, peak {int(np.abs(codes).max())}/2047")
    print(f"  body {len(body)} bytes, re-encoded {len(re_body)}")
    print(f"  codec round-trip byte-identical: {exact}")
    expect = (count + 1) // 2 * 3
    print(f"  body length matches count:       {len(body) == expect}"
          f"  (count {count} -> {expect} bytes)")
    print(f"  count is even (as reference):     {count % 2 == 0}")
    return 0 if exact else 1


def main():
    p = argparse.ArgumentParser(description="Akai MPC60 SND converter.")
    sub = p.add_subparsers(dest="cmd", required=True)

    c = sub.add_parser("convert", help="WAV folder (recursive) -> SND")
    c.add_argument("folder")
    c.add_argument("-o", "--out", default="./snd_out")
    c.add_argument("--name", default="KIT")
    c.add_argument("--rate", type=int, default=MPC60_RATE)
    c.add_argument("--no-normalize", action="store_true",
                   help="do not peak-normalise (wastes 12-bit headroom)")
    c.add_argument("--no-clamp", action="store_true",
                   help="copy unknown header pointers verbatim instead of "
                        "clamping them to the new sample count")
    c.add_argument("--strip", help="drop this leading text from each name "
                   "(e.g. --strip Miami_ so the 8.3 filename isn't eaten by it)")
    c.add_argument("--prefix", help="prepend this to each name")
    c.add_argument("--snd-template", help="donor .SND for the header")
    c.add_argument("--pgm-template", help="donor .PGM to rename in place")
    c.set_defaults(func=cmd_convert)

    d = sub.add_parser("decode", help="SND -> WAV")
    d.add_argument("file")
    d.add_argument("-o", "--out")
    d.add_argument("--rate", type=int, default=MPC60_RATE)
    d.set_defaults(func=cmd_decode)

    i = sub.add_parser("inspect", help="dump and check a SND header")
    i.add_argument("file")
    i.set_defaults(func=cmd_inspect)

    v = sub.add_parser("verify", help="prove the codec round-trips a real file")
    v.add_argument("file")
    v.set_defaults(func=cmd_verify)

    args = p.parse_args()
    sys.exit(args.func(args) or 0)


if __name__ == "__main__":
    main()