#!/usr/bin/env python3
"""
mpc60drive - inspect and edit an MPC60 ZuluSCSI drive image (.hda).

An MPC60 drive is just N FAT12 partitions laid end to end with NO partition
table. The MPC letters them A, B, C ... Z. Each one is ~29 MB because FAT12
with 8 KB clusters cannot address more.

The edit loop needs no special tools:

    ./mpc60drive.py list    drive.hda
    ./mpc60drive.py extract drive.hda C -o C.img
    open C.img                 # macOS mounts it - drag .SND/.PGM in Finder
    ./mpc60drive.py replace drive.hda C C.img

Also:
    ./mpc60drive.py init   new.hda            # blank 26-partition drive
    ./mpc60drive.py format drive.hda H        # wipe one partition

REMEMBER THE REAL LIMIT: the MPC60 has 750 KB of sound memory (1.5 MB if
expanded) - about 12.8 seconds of audio. A 29 MB partition holds far more
than the machine can load at once, so keep each PROGRAM under ~700 KB.
`list` shows the KB per partition so you can see when you are over.
"""
import argparse, os, struct, sys

SECT = 512
RESERVED, NFATS, SPF, ROOT_ENTS = 1, 2, 11, 224
ROOT_SECT = ROOT_ENTS * 32 // SECT
DATA_START = RESERVED + NFATS * SPF + ROOT_SECT
SPC = 16


def geometry(path):
    """Read the partition size out of the drive's own first boot sector."""
    with open(path, "rb") as f:
        boot = f.read(SECT)
    if boot[3:11] != b"MPC-60  ":
        print(f"  warning: OEM is {boot[3:11]!r}, expected b'MPC-60  '", file=sys.stderr)
    tot = struct.unpack("<H", boot[19:21])[0]
    if not tot:
        sys.exit("partition A has no BPB - is this really an MPC60 drive image?")
    return boot, tot, tot * SECT


def npartitions(path, psize):
    return os.path.getsize(path) // psize


def idx(letter):
    L = letter.strip().upper()
    if len(L) != 1 or not ("A" <= L <= "Z"):
        sys.exit(f"partition must be a single letter A-Z, got {letter!r}")
    return ord(L) - 65


def read_file(blob, start, size):
    """Walk the FAT12 chain and return a file's bytes."""
    raw = blob[RESERVED * SECT:][:SPF * SECT]
    fat = []
    for i in range(0, len(raw) - 2, 3):
        b0, b1, b2 = raw[i], raw[i + 1], raw[i + 2]
        fat.append(b0 | ((b1 & 0x0F) << 8))
        fat.append((b1 >> 4) | (b2 << 4))
    out, c, guard = b"", start, 0
    while 2 <= c < len(fat) and fat[c] < 0xFF8 or (c == start and guard == 0):
        o = (DATA_START + (c - 2) * SPC) * SECT
        out += blob[o:o + SPC * SECT]
        if len(out) >= size:
            break
        nxt = fat[c] if c < len(fat) else 0xFFF
        if nxt >= 0xFF8:
            break
        c = nxt
        guard += 1
        if guard > 4000:
            break
    return out[:size]


def program_sizes(blob):
    """-> [(program name, bytes of sound it references)] - the RAM figure."""
    root = blob[(RESERVED + NFATS * SPF) * SECT:][:ROOT_ENTS * 32]
    snds, pgms = {}, []
    for j in range(ROOT_ENTS):
        e = root[j * 32:j * 32 + 32]
        if e[0] in (0x00, 0xE5) or (e[11] & 0x18):
            continue
        nm = e[0:11].decode("ascii", "replace")
        start = int.from_bytes(e[26:28], "little")
        size = int.from_bytes(e[28:32], "little")
        if nm[8:11] == "PGM":
            pgms.append((nm[:8].strip(), start, size))
        else:
            hdr = read_file(blob, start, 39)
            if len(hdr) >= 18:
                snds[hdr[2:18].decode("ascii", "replace").rstrip()] = max(0, size - 39)
    out = []
    for name, start, size in pgms:
        data = read_file(blob, start, size)
        total, off, empt = 0, 0x24, 0
        while off + 16 <= len(data) and empt < 8:
            if data[off] not in (0x00, 0x20):
                total += snds.get(data[off:off + 16].decode("ascii", "replace").rstrip(), 0)
                empt = 0
            else:
                empt += 1
            off += 17
        out.append((name, total))
    return out


def scan_partition(blob):
    """-> (n_sounds, n_programs, sound_bytes, [program names])"""
    root = blob[(RESERVED + NFATS * SPF) * SECT:][:ROOT_ENTS * 32]
    snd = pgm = 0
    total = 0
    names = []
    for j in range(ROOT_ENTS):
        e = root[j * 32:j * 32 + 32]
        if e[0] in (0x00, 0xE5) or (e[11] & 0x18):
            continue
        name = e[0:11].decode("ascii", "replace")
        size = int.from_bytes(e[28:32], "little")
        if name[8:11] == "PGM":
            pgm += 1
            names.append(name[:8].strip())
        else:
            snd += 1
            total += max(0, size - 39)
    return snd, pgm, total, names


def blank_partition(boot, tot):
    """A validly formatted, empty FAT12 volume using this drive's own BPB."""
    img = bytearray(tot * SECT)
    img[0:SECT] = boot
    fat = bytearray(SPF * SECT)
    fat[0:3] = bytes([0xF8, 0xFF, 0xFF])          # cluster 0 and 1 sentinels
    for n in range(NFATS):
        o = (RESERVED + n * SPF) * SECT
        img[o:o + SPF * SECT] = fat
    return bytes(img)


def cmd_list(a):
    boot, tot, psize = geometry(a.drive)
    n = npartitions(a.drive, psize)
    print(f"{a.drive}: {os.path.getsize(a.drive)} bytes, {n} partitions of "
          f"{psize} bytes ({tot} sectors each)\n")
    print(f"  {'':3} {'sounds':>7} {'pgms':>5} {'disk KB':>9} {'worst pgm':>10}  programs")
    budget = 700
    with open(a.drive, "rb") as f:
        for i in range(n):
            f.seek(i * psize)
            blob = f.read(psize)
            snd, pgm, tb, names = scan_partition(blob)
            progs = program_sizes(blob) if pgm else []
            over = [n for n, b in progs if b > budget * 1024]
            worst = max((b for _, b in progs), default=0) / 1024
            flag = f"   <- {len(over)} program(s) OVER: {', '.join(over[:3])}" if over else ""
            shown = ", ".join(names[:4]) + (" ..." if len(names) > 4 else "")
            print(f"  {chr(65+i)}   {snd:>7} {pgm:>5} {tb/1024:>9.0f} {worst:>10.0f}  "
                  f"{shown or '-'}{flag}")
    print(f"\n  MPC60 sound memory: 750 KB stock / 1500 KB expanded "
          f"(~12.8 / 25.6 seconds)")


def cmd_extract(a):
    boot, tot, psize = geometry(a.drive)
    i = idx(a.letter)
    if i >= npartitions(a.drive, psize):
        sys.exit(f"drive only has {npartitions(a.drive, psize)} partitions")
    out = a.out or f"{a.letter.upper()}.img"
    with open(a.drive, "rb") as f:
        f.seek(i * psize)
        data = f.read(psize)
    open(out, "wb").write(data)
    snd, pgm, tb, _ = scan_partition(data)
    print(f"  {a.letter.upper()} -> {out}  ({len(data)} bytes, {snd} sounds, {pgm} programs)")
    print(f"  now:  open {out}      # macOS mounts it; drag files in Finder")
    print(f"  then: {sys.argv[0]} replace {a.drive} {a.letter.upper()} {out}")


def cmd_replace(a):
    boot, tot, psize = geometry(a.drive)
    i = idx(a.letter)
    if i >= npartitions(a.drive, psize):
        sys.exit(f"drive only has {npartitions(a.drive, psize)} partitions")
    data = open(a.image, "rb").read()
    if len(data) != psize:
        sys.exit(f"{a.image} is {len(data)} bytes, partition is {psize}. "
                 f"Refusing - a wrong-sized write would corrupt the next partition.")
    if data[3:11] != b"MPC-60  ":
        sys.exit(f"{a.image} has OEM {data[3:11]!r}; refusing to write a "
                 f"non-MPC60 volume. Extract a partition first and edit that.")
    snd, pgm, tb, _ = scan_partition(data)
    if tb > 700 * 1024:
        print(f"  note: {tb/1024:.0f} KB of sound here, more than the MPC60 can "
              f"hold at once. Fine on disk; keep each PROGRAM under ~700 KB.")
    with open(a.drive, "r+b") as f:
        f.seek(i * psize)
        f.write(data)
    print(f"  {a.image} -> {a.drive} partition {a.letter.upper()}  "
          f"({snd} sounds, {pgm} programs)")


def cmd_init(a):
    if a.boot:
        boot = open(a.boot, "rb").read(SECT)
        tot = struct.unpack("<H", boot[19:21])[0]
    else:
        sys.exit("need --boot: a 512-byte MPC60 boot sector, or the first "
                 "512 bytes of a drive the MPC formatted itself")
    blank = blank_partition(boot, tot)
    with open(a.out, "wb") as f:
        for _ in range(a.partitions):
            f.write(blank)
    print(f"  {a.out}: {a.partitions} blank partitions of {len(blank)} bytes "
          f"= {os.path.getsize(a.out)} bytes")


def cmd_format(a):
    boot, tot, psize = geometry(a.drive)
    i = idx(a.letter)
    with open(a.drive, "r+b") as f:
        f.seek(i * psize)
        f.write(blank_partition(boot, tot))
    print(f"  partition {a.letter.upper()} wiped to an empty formatted volume")


def main():
    p = argparse.ArgumentParser(description="Inspect and edit MPC60 drive images.")
    s = p.add_subparsers(dest="cmd", required=True)

    c = s.add_parser("list", help="what is in every partition")
    c.add_argument("drive"); c.set_defaults(func=cmd_list)

    c = s.add_parser("extract", help="pull one partition out as a mountable .img")
    c.add_argument("drive"); c.add_argument("letter"); c.add_argument("-o", "--out")
    c.set_defaults(func=cmd_extract)

    c = s.add_parser("replace", help="write an edited .img back into a partition")
    c.add_argument("drive"); c.add_argument("letter"); c.add_argument("image")
    c.set_defaults(func=cmd_replace)

    c = s.add_parser("init", help="create a fresh blank multi-partition drive")
    c.add_argument("out"); c.add_argument("--partitions", type=int, default=26)
    c.add_argument("--boot", help="512-byte MPC60 boot sector to seed from")
    c.set_defaults(func=cmd_init)

    c = s.add_parser("format", help="wipe one partition")
    c.add_argument("drive"); c.add_argument("letter"); c.set_defaults(func=cmd_format)

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


if __name__ == "__main__":
    main()