#!/usr/bin/env python3
"""
linn9000build - turn a folder of WAV files into a Linn 9000 disk image.

    ./linn9000build.py MYKIT/ -o mykit.img
    ./img2hfe.py mykit.img -o mykit.hfe      # then copy to the Gotek stick

The Linn is small and strict, so the tool tells you what it dropped and why
rather than silently truncating:

    64 samples maximum          one directory slot each
    65536 bytes per sample      the block list is 16 bytes wide, so 2.05 s
    716 KB total                179 blocks of 4096 bytes, about 22.9 s
    6-character names           upper case, letters digits _ #

Audio is resampled to 32 kHz mono and mu-law encoded, which is what the
machine expects.  Anything longer than 2.05 s is cut, with a warning.
"""
import argparse, glob, os, struct, sys, wave
import numpy as np
import linnfmt as L


def load_wav(path):
    """-> float array in -1..1, sample rate.  Mono-ed by averaging."""
    try:
        with wave.open(path, "rb") as w:
            n, ch, sw, sr = (w.getnframes(), w.getnchannels(),
                             w.getsampwidth(), w.getframerate())
            raw = w.readframes(n)
    except (wave.Error, EOFError) as e:
        raise SystemExit(f"  cannot read {path}: {e}")
    if sw == 1:                                   # 8-bit WAV is unsigned
        a = np.frombuffer(raw, np.uint8).astype(np.float32) - 128.0
        a /= 128.0
    elif sw == 2:
        a = np.frombuffer(raw, "<i2").astype(np.float32) / 32768.0
    elif sw == 3:
        b = np.frombuffer(raw, np.uint8).reshape(-1, 3).astype(np.int32)
        v = (b[:, 0] | (b[:, 1] << 8) | (b[:, 2] << 16))
        v = np.where(v & 0x800000, v - 0x1000000, v)
        a = v.astype(np.float32) / 8388608.0
    elif sw == 4:
        a = np.frombuffer(raw, "<i4").astype(np.float32) / 2147483648.0
    else:
        raise SystemExit(f"  {path}: unsupported sample width {sw*8}-bit")
    if ch > 1:
        a = a[:len(a) // ch * ch].reshape(-1, ch).mean(axis=1)
    return a, sr


def resample(x, sr_in, sr_out):
    """Band-limited resample through the frequency domain."""
    if sr_in == sr_out or len(x) == 0:
        return x
    n_out = max(1, int(round(len(x) * sr_out / sr_in)))
    X = np.fft.rfft(x)
    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))


def encode(x):
    """float -1..1 -> mu-law bytes, trailing silence trimmed."""
    q = np.clip(x * 32767.0, -32635, 32635)
    b = bytes(L.ulaw_encode(v) for v in q)
    end = len(b)
    while end and b[end - 1] in (0x00, 0x80):
        end -= 1
    return b[:end]


def build(items, verbose=True):
    """items: [(name, mu-law bytes)] -> a 737280-byte image."""
    img = bytearray(bytes([L.EMPTY]) * L.DISK)
    nxt, placed = 1, []
    for name, audio in items:
        total = L.HDR + len(audio)
        nblk = -(-total // L.BLOCK)
        if nxt + nblk - 1 > L.MAX_BLOCKS:
            print(f"  full: no room for {name}, skipped", file=sys.stderr)
            continue
        blocks = list(range(nxt, nxt + nblk))
        nxt += nblk
        rec = (struct.pack("<H", len(audio)) + name.encode("ascii")
               + L.HDR_TAIL + audio)
        rec += b"\x00" * (nblk * L.BLOCK - len(rec))
        for i, b in enumerate(blocks):
            img[b * L.BLOCK:(b + 1) * L.BLOCK] = rec[i * L.BLOCK:(i + 1) * L.BLOCK]
        placed.append((name, total, blocks))
    for i, (name, total, blocks) in enumerate(placed):
        e = bytearray(b"\x00" * L.ENTRY)
        e[0:6] = name.encode("ascii")
        e[6:10] = b"\xe5\xe5\xe5\xe5"
        e[10], e[11] = L.DIR_FLAG
        struct.pack_into("<H", e, 13, total)
        e[16:16 + len(blocks)] = bytes(blocks)
        img[i * L.ENTRY:(i + 1) * L.ENTRY] = e
    return bytes(img), placed


def main():
    p = argparse.ArgumentParser(description=__doc__,
            formatter_class=argparse.RawDescriptionHelpFormatter)
    p.add_argument("src", help="folder of WAVs, or a single WAV")
    p.add_argument("-o", "--out", required=True)
    p.add_argument("--gain", choices=("kit", "sample", "none"), default="kit",
                   help="kit: one gain for the whole folder, preserving the "
                        "balance between sounds (default). sample: normalise "
                        "each on its own, which flattens the kit. none.")
    p.add_argument("--headroom", type=float, default=1.0,
                   help="peak to normalise to, default 1.0")
    a = p.parse_args()

    files = (sorted(glob.glob(os.path.join(glob.escape(a.src), "*.[wW][aA][vV]")))
             if os.path.isdir(a.src) else [a.src])
    if not files:
        raise SystemExit(f"  no WAVs in {a.src}")

    loaded = []
    for f in files:
        x, sr = load_wav(f)
        x = resample(x, sr, L.RATE)
        loaded.append((os.path.splitext(os.path.basename(f))[0], x, sr))

    peak = max((np.abs(x).max() for _, x, _ in loaded if len(x)), default=1.0)
    items, used = [], set()
    print(f"\n  {len(loaded)} files -> {a.out}\n")
    for base, x, sr in loaded[:L.DIR_SLOTS]:
        if a.gain == "kit" and peak > 0:
            x = x * (a.headroom / peak)
        elif a.gain == "sample" and len(x) and np.abs(x).max() > 0:
            x = x * (a.headroom / np.abs(x).max())
        au = encode(x)
        cut = ""
        if len(au) > L.MAX_SAMPLE - L.HDR:
            au = au[:L.MAX_SAMPLE - L.HDR]
            cut = f"  CUT to {L.MAX_SAMPLE/L.RATE:.2f}s"
        name = L.clean_name(base)
        while name in used:                       # names must stay unique
            name = L.clean_name(base[:5] + str(len(used)))
        used.add(name)
        items.append((name, au))
        print(f"  {base[:22]:<22} {sr:>6} Hz -> {name}  "
              f"{len(au):>6} B  {len(au)/L.RATE:>5.2f}s{cut}")
    if len(loaded) > L.DIR_SLOTS:
        print(f"\n  only the first {L.DIR_SLOTS} files fit the directory; "
              f"{len(loaded)-L.DIR_SLOTS} ignored", file=sys.stderr)

    img, placed = build(items)
    open(a.out, "wb").write(img)
    used_blocks = sum(len(b) for _, _, b in placed)
    print(f"\n  {len(placed)} samples, {used_blocks}/{L.MAX_BLOCKS} blocks "
          f"({used_blocks*L.BLOCK/1024:.0f} KB), {L.DISK} bytes -> {a.out}\n")


if __name__ == "__main__":
    main()