#!/usr/bin/env python3
# nsi2hfe.py
#
# Convert NSI file to HFEv3. v0.1
#
# Written & released by Keir Fraser <keir.xen@gmail.com> and Eric Anderson
# <ejona86@gmail.com>
#
# This is free and unencumbered software released into the public domain.
# See the file COPYING for more details, or visit <http://unlicense.org>.

import sys,struct,argparse

def mfm(bs):
  os = bytearray()
  last_d = False
  o = 0
  for b in bs:
    for _ in range(2):
      for _ in range(4):
        o >>= 2
        d = bool(b & 0x80)
        o |= (not (d or last_d)) << 6
        o |= d << 7
        last_d = d
        b <<= 1
      os.append(o)
  return os

# Encode fm with twice as many bits to avoid runs of 1s that could be
# interpreted as opcodes
def fm(bs):
  os = bytearray()
  o = 0
  for b in bs:
    for _ in range(4):
      for _ in range(2):
        o >>= 4
        d = bool(b & 0x80)
        o |= 1 << 4
        o |= d << 6
        b <<= 1
      os.append(o)
  return os

def check(bs):
  c = 0
  for b in bs:
    c ^= b
    c = ((c & 0x7F) << 1) | (c >> 7)
  return bytes([c])

def main(argv):
  parser = argparse.ArgumentParser(
    formatter_class=argparse.ArgumentDefaultsHelpFormatter)
  parser.add_argument("--rate", type=int, default=250,
                      help="data rate, kbit/s")
  parser.add_argument("--rpm", type=int, default=300,
                      help="rotational rate, rpm")
  parser.add_argument("--cyls", type=int, default=35,
                      help="number of cylinders")
  parser.add_argument("--sides", type=int, default=-1,
                      help="number of sides (-1=auto)")
  parser.add_argument("--hard-sectors", type=int, default=10,
                      help="number of hard sectors")
  parser.add_argument("--infile-sector-size", type=int, default=-1,
                      help="bytes per sector (-1=auto)")
  parser.add_argument("infile", help="input filename")
  parser.add_argument("outfile", help="output filename")
  args = parser.parse_args(argv[1:])

  in_f = open(args.infile, "rb")
  size = in_f.seek(0, 2)
  in_f.seek(0, 0)

  if size not in (89600, 179200, 358400):
    print("Unexpected file size: %d, size")
    return 1

  if args.sides == -1:
    args.sides = 2 if size == 358400 else 1
  if args.infile_sector_size == -1:
    args.infile_sector_size = 256 if size == 89600 else 512
  use_mfm = args.infile_sector_size == 512
  if use_mfm:
    preamble = bytes(32) + b'\xfb' * 2
    filler = b'\x55'
  else:
    preamble = bytes(16) + b'\xfb'
    filler = b'\x11'

  v3 = args.hard_sectors != 0

  bits = (args.rate * 1000 * 60) // args.rpm
  bits *= 2 # clock bits
  bits *= 2 # 2 sides
  byte_cnt = (bits + 7) // 8 # convert to bytes, rounded up
  if not v3:
    byte_cnt = (byte_cnt + 15) & ~15 # round up to 16-byte boundary
  raw_bytes = byte_cnt
  if args.hard_sectors:
    byte_cnt += (args.hard_sectors + 1) * 2 # 2 sides
  blocks = (byte_cnt + 511) // 512 # convert to 512-byte blocks, rounded up

  print("Geometry: %u cylinders, %u sides" % (args.cyls, args.sides))
  print("%ukbit/s @ %uRPM -> %u Encoded Bits"
        % (args.rate, args.rpm, bits/2))
  print("Data per HFE Track: %u bytes, %u blocks" % (byte_cnt, blocks))

  # Header
  out_f = open(args.outfile, "wb")
  sig = b'HXCHFEV3' if v3 else b"HXCPICFE"
  out_f.write(struct.pack("<8s4B2H2BH",
                          sig,        # signature
                          0,          # revision
                          args.cyls,  # nr_tracks
                          args.sides, # nr_sides
                          0xff,       # track_encoding
                          args.rate,  # bitrate
                          0,          # rpm
                          0xfe,       # interface_mode
                          1,          # rsvd
                          1))         # track_list_offset
  out_f.write(bytearray(b'\xff'*(512-20)))

  # TLUT
  tlut_blocks = (args.cyls*4 + 511) // 512
  base = 1 + tlut_blocks
  for i in range(args.cyls):
    out_f.write(struct.pack("<2H", base, byte_cnt))
    base += blocks
  out_f.write(bytearray(b'\xff'*(tlut_blocks*512-args.cyls*4)))

  # Data
  if not args.hard_sectors:
    out_f.write(bytearray(b'\x88'*(blocks*512*args.cyls)))
  else:
    for cyl in range(args.cyls):
      side = []
      for ss in range(2):
        trk_side = bytearray(b'\x88')*(raw_bytes//2)
        sector_len = float(raw_bytes//2) / args.hard_sectors
        if args.sides == 2 or ss == 0:
          if ss == 0:
              in_base = cyl * args.infile_sector_size * args.hard_sectors
          else:
              in_base = (2*args.cyls - cyl - 1) * args.infile_sector_size * args.hard_sectors
          last_sector_pos = len(trk_side)
          for sector in range(args.hard_sectors-1, -1, -1):
            in_f.seek(in_base + sector * args.infile_sector_size)
            raw_b = in_f.read(args.infile_sector_size)
            raw_b = preamble + raw_b + check(raw_b) + bytes(1)
            sector_pos = int(sector_len*sector)
            if use_mfm:
              raw_mfm = mfm(raw_b)
            else:
              raw_mfm = fm(raw_b)
            raw_mfm += filler * (last_sector_pos-sector_pos-len(raw_mfm))
            trk_side[sector_pos:sector_pos+len(raw_mfm)] = raw_mfm
            last_sector_pos = sector_pos
        trk_side.insert(-int(sector_len)//2, 0x8F)
        for sector in range(args.hard_sectors-1, -1, -1):
          trk_side.insert(int(sector_len*sector), 0x8F)
        assert len(trk_side) == byte_cnt//2
        trk_side += b'\x0F'*(256 - len(trk_side)%256)
        assert len(trk_side) == blocks*256
        side.append(trk_side)
      if args.sides == 1:
        side.append(trk_side)

      trk = bytearray()
      for pos in range(0, len(side[0]), 256):
        trk.extend(side[0][pos:pos+256])
        trk.extend(side[1][pos:pos+256])
      out_f.write(trk)
  return 0
    
if __name__ == "__main__":
    sys.exit(main(sys.argv))