#!/usr/bin/env python3
"""
img2hfe - write a raw 720K image into an HFE file for a Gotek / FlashFloppy.

    ./img2hfe.py mykit.img -o mykit.hfe

Rather than synthesising MFM tracks from nothing, this patches a real
formatted blank disk supplied by the drive vendor, replacing each sector's
512 data bytes and recomputing its CRC.  Every gap, sync pattern, address
mark and timing detail therefore stays exactly as a working Linn disk has
them, and only the payload changes.

MFM rule: each data bit becomes two cells, a clock and the bit.  The clock
is 1 only when the previous data bit and this one are both 0.  Changing a
data byte therefore also changes the clock of the bit that follows it, so
the first cell after the CRC is recomputed too.
"""
import argparse, os, struct, sys

SYNC = 0x4489


def crc16(data, crc=0xFFFF):
    for b in data:
        crc ^= b << 8
        for _ in range(8):
            crc = ((crc << 1) ^ 0x1021) & 0xFFFF if crc & 0x8000 else (crc << 1) & 0xFFFF
    return crc


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


def from_bits(bits):
    out = bytearray(len(bits) // 8)
    for i, b in enumerate(bits):
        if b:
            out[i >> 3] |= 1 << (i & 7)
    return bytes(out)


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


def dec_byte(cells, i):
    v = 0
    for k in range(8):
        v = (v << 1) | cells[i + 2 * k + 1]
    return v


def enc_bytes(data, prev_bit):
    """bytes -> MFM cells, given the data bit preceding them."""
    cells = bytearray()
    p = prev_bit
    for byte in data:
        for k in range(7, -1, -1):
            d = (byte >> k) & 1
            cells.append(1 if (p == 0 and d == 0) else 0)
            cells.append(d)
            p = d
    return cells, p


def find_fields(cells):
    """-> [(cyl, head, sec, bit index where the 512 data bytes start)]"""
    out, n, i, reg, pending = [], len(cells), 0, 0, None
    while i < n - 16:
        reg = ((reg << 1) | cells[i]) & 0xFFFF
        if reg == SYNC:
            j, k = i + 1, 1
            while j + 16 <= n and peek(cells, j) == SYNC:
                k += 1; j += 16
            if k >= 3 and j + 16 <= n:
                mark = dec_byte(cells, j); j += 16
                if mark == 0xFE and j + 64 <= n:
                    pending = (dec_byte(cells, j), dec_byte(cells, j + 16),
                               dec_byte(cells, j + 32))
                    i = j + 64
                    continue
                if mark == 0xFB and pending and j + 514 * 16 <= n:
                    out.append(pending + (j,))
                    pending = None
                    i = j + 514 * 16
                    continue
            i = j
            continue
        i += 1
    return out


def main():
    here = os.path.dirname(os.path.abspath(__file__))
    p = argparse.ArgumentParser(description=__doc__,
            formatter_class=argparse.RawDescriptionHelpFormatter)
    p.add_argument("img")
    p.add_argument("-o", "--out", required=True)
    p.add_argument("--template", default=os.path.join(
        here, "..", "reference", "blank-linn9000-720k.hfe"))
    p.add_argument("--spt", type=int, default=9)
    a = p.parse_args()

    img = open(a.img, "rb").read()
    tpl = bytearray(open(a.template, "rb").read())
    if tpl[:8] != b"HXCPICFE":
        sys.exit("template is not an HFE v1 file")
    ntrack, nside = tpl[9], tpl[10]
    lut = struct.unpack_from("<H", tpl, 18)[0] * 512
    print(f"  template {os.path.basename(a.template)}: {ntrack} tracks, "
          f"{nside} sides")

    written = 0
    for c in range(ntrack):
        off, ln = struct.unpack_from("<HH", tpl, lut + c * 4)
        raw = tpl[off * 512: off * 512 + ln]
        side = [bytearray(), bytearray()]
        idx = []
        for i in range(0, len(raw), 256):
            s = (i // 256) & 1
            idx.append((s, len(side[s]), min(256, len(raw) - i)))
            side[s] += raw[i:i + 256]
        for h in range(nside):
            cells = to_bits(side[h])
            for (cc, hh, ss, bit) in find_fields(cells):
                lba = (c * nside + h) * a.spt + (ss - 1)
                data = img[lba * 512:(lba + 1) * 512]
                if len(data) != 512:
                    continue
                blk = bytes(data) + struct.pack(">H",
                        crc16(b"\xa1\xa1\xa1\xfb" + bytes(data)))
                new, last = enc_bytes(blk, 1)      # 0xFB ends in a 1 bit
                cells[bit:bit + len(new)] = new
                nxt = bit + len(new)               # fix the following clock
                if nxt + 1 < len(cells):
                    cells[nxt] = 1 if (last == 0 and cells[nxt + 1] == 0) else 0
                written += 1
            side[h] = bytearray(from_bits(cells))
        merged = bytearray()
        for s, o, ln2 in idx:
            merged += side[s][o:o + ln2]
        tpl[off * 512: off * 512 + ln] = merged

    open(a.out, "wb").write(bytes(tpl))
    print(f"  patched {written} sectors -> {a.out} ({len(tpl)} bytes)")


if __name__ == "__main__":
    main()