#!/usr/bin/env python3
"""xdftool - read/write Human68k FAT12 floppy images (.XDF) from Linux.

Why this exists: mtools cannot handle these images. Human68k 2HD floppies use
1024-byte logical sectors, and mformat mis-sizes the volume while mcopy fails
with rc=1 regardless of configuration. imgtool has no Human68k module. So the
JoeyLib X68000 gate needs its own reader/writer to get a test binary onto a
disk and its log back off again.

The filesystem itself is ordinary FAT12; only the sector size is unusual, and
every offset in the BPB is already expressed in sectors, so nothing needs
special-casing beyond reading bytesPerSector rather than assuming 512.

This does NOT create bootable images. Mint one once by running Human68k's own
FORMAT.X/SYS.X inside an emulator (they drive the FDD through IOCS), keep it as
a template, and copy it per run -- the same pattern the IIgs gate uses with
gsos-system.po.

Usage:
    xdftool.py list    <image>
    xdftool.py add     <image> <hostfile> [name]
    xdftool.py delete  <image> <name>
    xdftool.py extract <image> <name> <hostfile>
    xdftool.py free    <image>
"""

import os
import struct
import sys

DIR_ENTRY_SIZE = 32
ATTR_VOLUME    = 0x08
ATTR_DIR       = 0x10
ATTR_ARCHIVE   = 0x20
FREE_MARKER    = 0x00
DELETED_MARKER = 0xE5
EOC_MIN        = 0xFF8   # >= this in a FAT12 entry means end-of-chain


class Xdf:
    def __init__(self, path, writable=False):
        self.path = path
        self.writable = writable
        with open(path, "rb") as fp:
            self.data = bytearray(fp.read())
        self._parse_bpb()

    def _parse_bpb(self):
        (self.bps, self.spc, self.reserved, self.nfats, self.root_entries,
         self.total_sectors, self.media, self.spf, self.spt, self.heads,
         self.hidden) = struct.unpack("<HBHBHHBHHHH", self.data[11:30])
        if self.bps == 0 or self.spc == 0:
            raise ValueError(f"{self.path}: implausible BPB (bytesPerSector={self.bps})")
        self.fat_start  = self.reserved * self.bps
        self.root_start = (self.reserved + self.nfats * self.spf) * self.bps
        self.data_start = self.root_start + self.root_entries * DIR_ENTRY_SIZE
        self.cluster_bytes = self.spc * self.bps
        # Clusters are numbered from 2, and the last valid one is bounded by
        # how much space is actually left after the root directory.
        self.max_cluster = 1 + (len(self.data) - self.data_start) // self.cluster_bytes

    # ----- FAT12 entry access -----

    def fat_get(self, cluster):
        off = self.fat_start + (cluster * 3) // 2
        pair = self.data[off] | (self.data[off + 1] << 8)
        return (pair >> 4) if (cluster & 1) else (pair & 0x0FFF)

    def fat_set(self, cluster, value):
        # Mirror into every FAT copy so the image stays self-consistent.
        for fat in range(self.nfats):
            base = self.fat_start + fat * self.spf * self.bps
            off = base + (cluster * 3) // 2
            pair = self.data[off] | (self.data[off + 1] << 8)
            if cluster & 1:
                pair = (pair & 0x000F) | ((value & 0x0FFF) << 4)
            else:
                pair = (pair & 0xF000) | (value & 0x0FFF)
            self.data[off] = pair & 0xFF
            self.data[off + 1] = (pair >> 8) & 0xFF

    def free_clusters(self):
        return [c for c in range(2, self.max_cluster + 1) if self.fat_get(c) == 0]

    def cluster_offset(self, cluster):
        return self.data_start + (cluster - 2) * self.cluster_bytes

    # ----- Directory -----

    def _entries(self):
        for i in range(self.root_entries):
            off = self.root_start + i * DIR_ENTRY_SIZE
            yield i, off, self.data[off:off + DIR_ENTRY_SIZE]

    def listdir(self):
        out = []
        for _, _, e in self._entries():
            if e[0] == FREE_MARKER:
                break
            if e[0] == DELETED_MARKER or (e[11] & (ATTR_VOLUME | ATTR_DIR)):
                continue
            name = e[0:8].decode("ascii", "replace").rstrip()
            ext = e[8:11].decode("ascii", "replace").rstrip()
            # Human68k EXTENDED names: chars 9-18 of the stem live in bytes 12-21 (see encode_name_ext).
            # Fold them back in so a listing shows the full name a guest wrote (e.g. the 9-char blob keys).
            more = e[12:22].split(b"\x00", 1)[0].decode("ascii", "replace").rstrip()
            stem = name + more
            out.append({
                "name": f"{stem}.{ext}" if ext else stem,
                "cluster": struct.unpack("<H", e[26:28])[0],
                "size": struct.unpack("<I", e[28:32])[0],
                "attr": e[11],
            })
        return out

    @staticmethod
    def encode_name(name):
        name = name.upper()
        stem, _, ext = name.partition(".")
        if len(stem) > 8 or len(ext) > 3:
            raise ValueError(f"'{name}' does not fit 8.3")
        return stem.ljust(8).encode("ascii") + ext.ljust(3).encode("ascii")

    @staticmethod
    def encode_name_ext(name):
        # Human68k EXTENDED names: the stem may run to 18 chars - the first 8 live in the classic
        # name field, chars 9-18 in the dir entry's bytes 12-21 (MS-DOS's reserved area).  Returns
        # (main11, ext10) for matching; ext10 is all-NUL for a plain 8.3 name.  RetroNet's blob names
        # ('A' + 8 hex = 9 chars) need this - without it a guest-written blob reads as "not in image".
        name = name.upper()
        stem, _, ext = name.partition(".")
        if len(stem) > 18 or len(ext) > 3:
            raise ValueError(f"'{name}' does not fit Human68k 18.3")
        main = stem[:8].ljust(8).encode("ascii") + ext.ljust(3).encode("ascii")
        extended = stem[8:].encode("ascii").ljust(10, b"\x00")
        return main, extended

    @staticmethod
    def entry_matches(e, main11, ext10):
        if bytes(e[0:11]).upper() != main11:
            return False
        got = bytes(e[12:22]).rstrip(b"\x00 ").upper()
        want = ext10.rstrip(b"\x00 ")
        return got == want

    def find(self, name):
        # Human68k PRESERVES filename case in the directory entry (unlike
        # MS-DOS, which upcases), so a program writing "joeylog.txt" leaves it
        # lowercase on disk. Match case-insensitively or extracting a
        # guest-written file fails with a confusing "not in image".
        main11, ext10 = self.encode_name_ext(name)
        for i, off, e in self._entries():
            if e[0] == FREE_MARKER:
                break
            if e[0] == DELETED_MARKER:
                continue
            if self.entry_matches(e, main11, ext10):
                return i, off, e
        return None, None, None

    # ----- File operations -----

    def find_in_subdir(self, dirname, leaf):
        # Resolve SUB/LEAF: a subdirectory's data clusters are themselves a directory table.  xdftool
        # only ever creates one-cluster subdirs, but the guest may have grown one, so follow the chain.
        _, _, de = self.find(dirname)
        if de is None or not (de[11] & ATTR_DIR):
            return None
        main11, ext10 = self.encode_name_ext(leaf)
        cluster = struct.unpack("<H", de[26:28])[0]
        guard = 0
        while 2 <= cluster < EOC_MIN:
            base = self.cluster_offset(cluster)
            for i in range(self.cluster_bytes // DIR_ENTRY_SIZE):
                e = self.data[base + i * DIR_ENTRY_SIZE:base + (i + 1) * DIR_ENTRY_SIZE]
                if e[0] == FREE_MARKER:
                    return None
                if e[0] == DELETED_MARKER or (e[11] & (ATTR_VOLUME | ATTR_DIR)):
                    continue
                if self.entry_matches(e, main11, ext10):
                    return e
            cluster = self.fat_get(cluster)
            guard += 1
            if guard > self.max_cluster:
                return None
        return None

    def read_file(self, name):
        if "/" in name:                          # SAVES/RN.CFG: resolve through the subdirectory
            sub, _, leaf = name.partition("/")
            e = self.find_in_subdir(sub, leaf)
        else:
            _, _, e = self.find(name)
        if e is None:
            raise FileNotFoundError(f"{name} not in {self.path}")
        size = struct.unpack("<I", e[28:32])[0]
        cluster = struct.unpack("<H", e[26:28])[0]
        out = bytearray()
        guard = 0
        while 2 <= cluster < EOC_MIN and len(out) < size:
            off = self.cluster_offset(cluster)
            out += self.data[off:off + self.cluster_bytes]
            cluster = self.fat_get(cluster)
            guard += 1
            if guard > self.max_cluster:
                raise ValueError(f"{name}: cluster chain loops")
        return bytes(out[:size])

    def delete(self, name):
        _, off, e = self.find(name)
        if e is None:
            return False
        cluster = struct.unpack("<H", e[26:28])[0]
        while 2 <= cluster < EOC_MIN:
            nxt = self.fat_get(cluster)
            self.fat_set(cluster, 0)
            cluster = nxt
        self.data[off] = DELETED_MARKER
        return True

    def _alloc_write_payload(self, payload):
        # Allocate a FAT chain for `payload`, write it (zero-padding the final cluster so stale bytes never
        # leak), and return the head cluster (0 for an empty payload).  Shared by add_file and
        # add_file_in_subdir so the cluster-writing logic lives in exactly one place.
        cb = self.cluster_bytes
        need = (len(payload) + cb - 1) // cb
        if need == 0:
            return 0
        free = self.free_clusters()
        if len(free) < need:
            raise OSError(f"{self.path}: need {need} clusters, {len(free)} free "
                          f"({len(free) * cb} bytes)")
        chain = free[:need]
        for idx, cluster in enumerate(chain):
            off = self.cluster_offset(cluster)
            chunk = payload[idx * cb:(idx + 1) * cb]
            self.data[off:off + len(chunk)] = chunk
            if len(chunk) < cb:
                self.data[off + len(chunk):off + cb] = bytes(cb - len(chunk))
            self.fat_set(cluster, 0xFFF if idx == need - 1 else chain[idx + 1])
        return chain[0]

    def _put_entry_in_dir(self, first_cluster, entry):
        # Write a 32-byte directory `entry` into the first FREE/DELETED slot of the directory whose data
        # begins at first_cluster, extending the cluster chain by one when every slot is already used.
        per = self.cluster_bytes // DIR_ENTRY_SIZE
        cluster = first_cluster
        last = cluster
        guard = 0
        while 2 <= cluster < EOC_MIN:
            base = self.cluster_offset(cluster)
            for i in range(per):
                slot = base + i * DIR_ENTRY_SIZE
                if self.data[slot] in (FREE_MARKER, DELETED_MARKER):
                    self.data[slot:slot + DIR_ENTRY_SIZE] = entry
                    return
            last = cluster
            cluster = self.fat_get(cluster)
            guard += 1
            if guard > self.max_cluster:
                raise OSError(f"{self.path}: directory cluster chain loops")
        # Every slot is used: append a fresh, all-free cluster and write into its first slot.
        free = self.free_clusters()
        if not free:
            raise OSError(f"{self.path}: no free cluster to extend directory")
        newc = free[0]
        self.fat_set(last, newc)
        self.fat_set(newc, 0xFFF)
        base = self.cluster_offset(newc)
        self.data[base:base + self.cluster_bytes] = bytes(self.cluster_bytes)
        self.data[base:base + DIR_ENTRY_SIZE] = entry

    def add_file(self, name, payload):
        self.delete(name)   # overwrite semantics
        # encode_name_ext writes the Human68k EXTENDED name (a strict superset of 8.3: ext10 is all-NUL
        # for a plain 8.3 name), so names up to 18 chars - like RetroNet's 9-char blob keys - land in a
        # form the guest and xdftool's own find()/read_file() both resolve.
        main11, ext10 = self.encode_name_ext(name)
        head = self._alloc_write_payload(payload)
        slot = None
        for _, off, e in self._entries():
            if e[0] in (FREE_MARKER, DELETED_MARKER):
                slot = off
                break
        if slot is None:
            raise OSError(f"{self.path}: root directory full ({self.root_entries} entries)")
        self.data[slot:slot + DIR_ENTRY_SIZE] = self._dir_entry(
            main11, ATTR_ARCHIVE, head, len(payload), ext10)

    def _dir_entry(self, name11, attr, cluster, size, ext10=None):
        entry = bytearray(DIR_ENTRY_SIZE)
        entry[0:11] = name11
        entry[11] = attr
        if ext10 is not None:                                                  # Human68k extended name
            entry[12:22] = ext10
        struct.pack_into("<H", entry, 22, (12 << 11))                          # 12:00:00
        struct.pack_into("<H", entry, 24, ((2026 - 1980) << 9) | (1 << 5) | 1)  # 2026-01-01
        struct.pack_into("<H", entry, 26, cluster)
        struct.pack_into("<I", entry, 28, size)
        return entry

    def add_file_in_subdir(self, dirname, name, payload):
        # Write `name` into subdirectory `dirname` (creating SUB/. and SUB/.. when the subdir is new), so a
        # guest reading e.g. SAVES/RN.CFG or SAVES/<blob> finds it.  If the subdir ALREADY exists (a second
        # file staged into the same SAVES/), add INTO it rather than creating a duplicate root entry the
        # guest's SUB lookup would never resolve.  Human68k FAT12 is otherwise ordinary FAT12.
        dstem = dirname.upper()
        if len(dstem) > 8:
            raise ValueError(f"subdir '{dirname}' does not fit 8.3")
        main11, ext10 = self.encode_name_ext(name)
        _, _, de = self.find(dirname)
        if de is not None and (de[11] & ATTR_DIR):
            head = self._alloc_write_payload(payload)
            entry = self._dir_entry(main11, ATTR_ARCHIVE, head, len(payload), ext10)
            self._put_entry_in_dir(struct.unpack("<H", de[26:28])[0], entry)
            return
        # Fresh subdirectory: reserve its cluster first so the payload allocation cannot reuse it, then
        # lay down '.', '..', and the file entry.
        free = self.free_clusters()
        if not free:
            raise OSError(f"{self.path}: no free cluster for subdirectory")
        dir_cluster = free[0]
        self.fat_set(dir_cluster, 0xFFF)
        head = self._alloc_write_payload(payload)
        cb = self.cluster_bytes
        doff = self.cluster_offset(dir_cluster)
        self.data[doff:doff + cb] = bytes(cb)
        self.data[doff:doff + 32] = self._dir_entry(b".          ", ATTR_DIR, dir_cluster, 0)
        self.data[doff + 32:doff + 64] = self._dir_entry(b"..         ", ATTR_DIR, 0, 0)
        self.data[doff + 64:doff + 96] = self._dir_entry(main11, ATTR_ARCHIVE, head, len(payload), ext10)
        slot = None
        for _, off, e in self._entries():
            if e[0] in (FREE_MARKER, DELETED_MARKER):
                slot = off
                break
        if slot is None:
            raise OSError(f"{self.path}: root directory full")
        self.data[slot:slot + DIR_ENTRY_SIZE] = self._dir_entry(
            dstem.ljust(8).encode("ascii") + b"   ", ATTR_DIR, dir_cluster, 0)

    def flush(self):
        if not self.writable:
            raise PermissionError("opened read-only")
        with open(self.path, "wb") as fp:
            fp.write(self.data)


def main(argv):
    if len(argv) < 3:
        print(__doc__.strip())
        return 2
    cmd, image = argv[1], argv[2]

    if cmd == "list":
        img = Xdf(image)
        print(f"{image}: {len(img.data)} bytes, {img.bps}-byte sectors, "
              f"{img.total_sectors} sectors, {img.spc} sec/cluster")
        for f in img.listdir():
            print(f"  {f['name']:<14} {f['size']:>9}  clus={f['cluster']}")
        return 0

    if cmd == "free":
        img = Xdf(image)
        n = len(img.free_clusters())
        print(f"{n * img.cluster_bytes} bytes free ({n} clusters)")
        return 0

    if cmd == "add":
        if len(argv) < 4:
            print("usage: xdftool.py add <image> <hostfile> [name]", file=sys.stderr)
            return 2
        host = argv[3]
        name = argv[4] if len(argv) > 4 else os.path.basename(host)
        img = Xdf(image, writable=True)
        with open(host, "rb") as fp:
            payload = fp.read()
        if "/" in name:                          # SAVES/RN.CFG -> a one-cluster subdirectory + the file
            sub, _, leaf = name.partition("/")
            img.add_file_in_subdir(sub, leaf, payload)
        else:
            img.add_file(name, payload)
        img.flush()
        print(f"added {name} ({os.path.getsize(host)} bytes) to {image}")
        return 0

    if cmd == "delete":
        if len(argv) < 4:
            print("usage: xdftool.py delete <image> <name>", file=sys.stderr)
            return 2
        img = Xdf(image, writable=True)
        if not img.delete(argv[3]):
            print(f"{argv[3]} not in {image}", file=sys.stderr)
            return 1
        img.flush()
        print(f"deleted {argv[3]} from {image}")
        return 0

    if cmd == "extract":
        if len(argv) < 5:
            print("usage: xdftool.py extract <image> <name> <hostfile>", file=sys.stderr)
            return 2
        img = Xdf(image)
        with open(argv[4], "wb") as fp:
            fp.write(img.read_file(argv[3]))
        print(f"extracted {argv[3]} -> {argv[4]}")
        return 0

    print(f"unknown command '{cmd}'", file=sys.stderr)
    return 2


if __name__ == "__main__":
    sys.exit(main(sys.argv))