#!/usr/bin/env python3
"""
linn9000 - read Linn 9000 / Forat floppy images and extract the samples.

    sudo dd if=/dev/rdiskN of=disk.img bs=512 conv=noerror,sync
    ./linn9000.py list    disk.img
    ./linn9000.py extract disk.img -o out [--rate 32000]
    ./linn9000.py info    disk.img

FORMAT, recovered from two independently written disks and verified by
decoding the audio.  720K disk, 737280 bytes, no filesystem.

  0x0000  directory: 64 slots of 32 bytes
            0-5    name, 6 ASCII chars, space padded
            6-9    E5 E5 E5 E5   present on every live entry
            10-11  51 E5 or 51 FF.  UNKNOWN - does not track size, block
                   count or position, and both values occur on one disk.
            12     always 00 on both disks
            13-14  length in BYTES, u16 little-endian
            16+    block numbers, one byte each, terminated by 00
          A slot whose first byte is E5 is free.

  data    4096-byte blocks.  Block N begins at N * 4096, so the directory
          occupies block 0 and the first sample block is block 1 at 0x1000.
          Each sample's FIRST block opens with a 16-byte header:
            0-1    length of the audio that follows, u16 LE
                   (always the directory size minus 16, on 51/51 samples)
            2-7    the name again, 6 ASCII chars
            8-15   UNKNOWN, varies per sample
          Later blocks of the same sample are pure audio, no header.

          Audio is 8-bit MU-LAW COMPANDED, giving roughly 14-bit range:
            bit  7    sign
            bits 6-4  exponent
            bits 3-0  mantissa
            value = ((mantissa << 3) + 0x84) << exponent, minus the 0x84
                    bias, negated when the sign bit is set
          This is G.711 mu-law but stored WITHOUT the bit inversion that
          telephony applies, so decode the byte directly.  0x00 and 0x80
          both give exactly 0, which is why samples end on 0x80 and are
          then padded with 0x00 - in this encoding that join is silent.

          Decoding it as linear PCM is what makes it sound distorted, and
          it is not subtle: on the sustained BAS808 tone, total harmonic
          distortion is 0.2% read as mu-law against 24.2% read linearly.
          Mu-law measured lower on all 11 tonal samples tested across both
          disks.  Reading it complemented, the telephony way, gives 159%.

  How the block layout was settled, since it is easy to get wrong:
    - block size:  for every sample on both disks the length field falls in
      (4096*(n-1), 4096*n] where n is the number of block bytes listed, and
      ceil(size/4096) == n for all 51 samples.  No other block size fits.
    - base offset: with offset = N*4096 every one of the 51 samples begins
      on a non-zero byte.  With 0x800+(N-1)*4096 they average ~1000 bytes of
      leading silence, i.e. that base is 2048 bytes early.
    - continuity:  across 101 block joins on the two disks the waveform
      discontinuity is a median 0.5-1.0x an ordinary sample-to-sample step,
      i.e. the seams are inaudible.

SAMPLE RATE is not stored on the disk.  The default of 32000 comes from the
808 kick on disk 3: its fundamental is 585 samples per cycle, and a TR-808
kick is tuned near 55 Hz (A1), which puts the rate at ~32.2 kHz.  Pitch
scales linearly, so if a sound plays a fifth low, multiply --rate by 1.5.
"""
import argparse, os, struct, sys

DIR_SLOTS, ENTRY, BLOCK, EMPTY, DISK = 64, 32, 4096, 0xE5, 737280


def entries(img):
    """-> [(name, nbytes, [blocks], flag)]"""
    out = []
    for i in range(DIR_SLOTS):
        e = img[i * ENTRY:(i + 1) * ENTRY]
        if len(e) < ENTRY or e[0] == EMPTY:
            continue
        name = e[0:6].decode("ascii", "replace").rstrip()
        if not name:
            continue
        nbytes = struct.unpack("<H", e[13:15])[0]
        term = e[16:ENTRY].find(0)
        blocks = list(e[16:16 + (ENTRY - 16 if term < 0 else term)])
        if not blocks or -(-nbytes // BLOCK) != len(blocks):
            print(f"  warning: {name} has {len(blocks)} blocks for "
                  f"{nbytes} bytes, skipping", file=sys.stderr)
            continue
        out.append((name, nbytes, blocks, (e[10], e[11])))
    return out


HDR = 16


def read_sample(img, blocks, nbytes, trim=True):
    """-> (audio bytes, header). Strips the 16-byte header, and by default
    the trailing silence padding, which is often over half the stored size."""
    raw = b"".join(img[b * BLOCK:(b + 1) * BLOCK] for b in blocks)[:nbytes]
    hdr, data = raw[:HDR], raw[HDR:]
    if trim:
        end = len(data)
        while end and data[end - 1] in (0x00, 0x80):
            end -= 1
        data = data[:end]
    return data, hdr


def ulaw(b):
    """One mu-law byte -> a signed 14-bit sample, range +/-32124."""
    t = (((b & 0x0F) << 3) + 0x84) << ((b & 0x70) >> 4)
    return 0x84 - t if b & 0x80 else t - 0x84


def wav(pcm8, rate):
    """8-bit mu-law mono -> 16-bit WAV, which everything plays."""
    pcm = b"".join(struct.pack("<h", ulaw(b)) for b in pcm8)
    return (b"RIFF" + struct.pack("<I", 36 + len(pcm)) + b"WAVEfmt " +
            struct.pack("<IHHIIHH", 16, 1, 1, rate, rate * 2, 2, 16) +
            b"data" + struct.pack("<I", len(pcm)) + pcm)


def load(path):
    d = open(path, "rb").read()
    if len(d) != DISK:
        print(f"  note: {len(d)} bytes, expected {DISK} for a 720K disk",
              file=sys.stderr)
    return d


def cmd_list(a):
    img = load(a.image)
    es = entries(img)
    print(f"\n  {a.image}\n")
    print(f"  {'name':<8} {'bytes':>7} {'sec':>6}  {'flag':>7}  blocks")
    for name, n, bl, fl in es:
        run = f"{bl[0]}-{bl[-1]}" if bl == list(range(bl[0], bl[-1] + 1)) \
              else ",".join(map(str, bl))
        print(f"  {name:<8} {n:>7} {n/a.rate:>6.2f}  "
              f"{fl[0]:02x} {fl[1]:02x}  {run}")
    used = sorted({b for _, _, bl, _ in es for b in bl})
    print(f"\n  {len(es)} samples, {len(used)} blocks used of "
          f"{len(img)//BLOCK - 1}, highest {max(used) if used else 0}")
    print(f"  {sum(n for _, n, _, _ in es)/1024:.0f} KB of audio, "
          f"{sum(n for _, n, _, _ in es)/a.rate:.1f} s at {a.rate} Hz\n")


def cmd_info(a):
    img = load(a.image)
    es = entries(img)
    used = {b for _, _, bl, _ in es for b in bl}
    total = len(img) // BLOCK - 1
    print(f"\n  {a.image}")
    print(f"  directory  block 0, {DIR_SLOTS} slots of {ENTRY} bytes, "
          f"{len(es)} in use")
    print(f"  data       block N at N*{BLOCK}, 8-bit mu-law")
    print(f"  capacity   {total} blocks, {total - len(used)} free "
          f"({(total-len(used))*BLOCK/1024:.0f} KB)")
    gaps = sorted(set(range(1, max(used) + 1)) - used) if used else []
    print(f"  allocation {'contiguous' if not gaps else f'gaps at {gaps}'}\n")


def cmd_extract(a):
    img = load(a.image)
    os.makedirs(a.out, exist_ok=True)
    n = 0
    for name, nbytes, blocks, _ in entries(img):
        data, hdr = read_sample(img, blocks, nbytes, trim=not a.no_trim)
        inner = hdr[2:8].decode("ascii", "replace").rstrip()
        if inner != name:
            print(f"  warning: {name} header says {inner!r}", file=sys.stderr)
        safe = "".join(c if c.isalnum() or c in "-_" else "_" for c in name)
        path = os.path.join(a.out, f"{safe}.wav")
        with open(path, "wb") as f:
            f.write(wav(data, a.rate))
        print(f"  {name:<8} {len(data):>7} bytes "
              f"({len(data)/a.rate:.2f}s) -> {path}")
        n += 1
    print(f"\n  {n} samples at {a.rate} Hz into {a.out}/\n")


def main():
    p = argparse.ArgumentParser(description=__doc__,
                                formatter_class=argparse.RawDescriptionHelpFormatter)
    p.add_argument("--rate", type=int, default=32000,
                   help="playback rate, not stored on disk (default 32000)")
    s = p.add_subparsers(dest="cmd", required=True)
    for name, fn in (("list", cmd_list), ("info", cmd_info),
                     ("extract", cmd_extract)):
        c = s.add_parser(name)
        c.add_argument("image")
        if name == "extract":
            c.add_argument("-o", "--out", default="linn_samples")
            c.add_argument("--no-trim", action="store_true",
                           help="keep the trailing silence padding")
        c.set_defaults(fn=fn)
    a = p.parse_args()
    a.fn(a)


if __name__ == "__main__":
    main()