#!/usr/bin/env python3
"""
hfe2img - HxC/FlashFloppy HFE v1 bitstream -> raw sector image.

    ./hfe2img.py disk.hfe -o disk.img

HFE stores the raw MFM cell stream a floppy controller would see, with the
two sides interleaved in 256-byte blocks and bits LSB-first.  To get sectors
back you have to decode MFM: find the A1 sync marks (which are written with a
deliberately illegal clock so they cannot occur in data), then take every
other cell as a data bit.
"""
import argparse, struct, sys

SYNC = 0x4489          # an A1 byte with the missing clock bit


def bits_lsb(buf):
    out = bytearray()
    for byte in buf:
        for i in range(8):
            out.append((byte >> i) & 1)
    return out


def track_streams(data, ntrack, lut_off):
    """-> {(cyl, head): bit list}"""
    out = {}
    for c in range(ntrack):
        off, ln = struct.unpack_from("<HH", data, lut_off + c * 4)
        raw = data[off * 512: off * 512 + ln]
        side = [bytearray(), bytearray()]
        for i in range(0, len(raw), 256):
            side[(i // 256) & 1] += raw[i:i + 256]
        for h in (0, 1):
            out[(c, h)] = bits_lsb(side[h])
    return out


def decode_byte(cells, i):
    """16 cells -> one byte; data bits are the odd cells."""
    v = 0
    for k in range(8):
        v = (v << 1) | cells[i + 2 * k + 1]
    return v


def sectors(cells):
    """-> {(cyl, head, sec): 512 bytes}"""
    found, n, i = {}, len(cells), 0
    reg = 0
    pending = None
    while i < n - 16:
        reg = ((reg << 1) | cells[i]) & 0xFFFFFFFF
        if (reg & 0xFFFF) == SYNC:
            # count consecutive A1 syncs, then read the mark byte
            j = i + 1
            k = 1
            while j + 16 <= n and _peek(cells, j) == SYNC:
                k += 1; j += 16
            if k >= 3 and j + 16 <= n:
                mark = decode_byte(cells, j)
                j += 16
                if mark == 0xFE and j + 4 * 16 <= n:
                    c = decode_byte(cells, j); h = decode_byte(cells, j + 16)
                    s = decode_byte(cells, j + 32)
                    pending = (c, h, s)
                    i = j + 64
                    continue
                if mark == 0xFB and pending and j + 512 * 16 <= n:
                    buf = bytes(decode_byte(cells, j + 16 * b)
                                for b in range(512))
                    found[pending] = buf
                    pending = None
                    i = j + 512 * 16
                    continue
            i = j
            continue
        i += 1
    return found


def _peek(cells, i):
    v = 0
    for k in range(16):
        v = (v << 1) | cells[i + k]
    return v


def main():
    p = argparse.ArgumentParser()
    p.add_argument("hfe"); p.add_argument("-o", "--out", required=True)
    p.add_argument("--spt", type=int, default=9)
    a = p.parse_args()
    d = open(a.hfe, "rb").read()
    if d[:8] != b"HXCPICFE":
        sys.exit("not an HFE v1 file")
    ntrack, nside = d[9], d[10]
    lut = struct.unpack_from("<H", d, 18)[0] * 512
    print(f"  {ntrack} tracks, {nside} sides, encoding {d[11]}, "
          f"{struct.unpack_from('<H', d, 12)[0]} kbit/s")

    allsec, geom = {}, {}
    for (c, h), cells in track_streams(d, ntrack, lut).items():
        for (cc, hh, ss), buf in sectors(cells).items():
            allsec[(c, h, ss)] = buf
            geom.setdefault("sec", set()).add(ss)
    spt = a.spt or (max(geom["sec"]) if geom.get("sec") else 9)
    lo = min(geom["sec"]) if geom.get("sec") else 1
    size = ntrack * nside * spt * 512
    img = bytearray(b"\x00" * size)
    got = 0
    for (c, h, s), buf in allsec.items():
        lba = (c * nside + h) * spt + (s - lo)
        if 0 <= lba * 512 < size:
            img[lba * 512:(lba + 1) * 512] = buf
            got += 1
    open(a.out, "wb").write(img)
    print(f"  sectors {sorted(geom.get('sec', []))[:12]}  recovered {got} "
          f"of {ntrack*nside*spt}  -> {a.out} ({size} bytes)")


if __name__ == "__main__":
    main()