#!/usr/bin/env python3
"""
imd_cpmtools_vfii.py

A small "cpmtools for IMD" tailored to the SD Systems VersaFloppy II CP/M 2.2
IMD format like:
  SBC548N_vector_ZCB_serialA_04_05.IMD

Commands:
  info    <image.imd>
  ls      <image.imd>
  get     <image.imd> <CPMFILE> [--out <path>]
  put     <image.imd> <local_path> [--as CPMNAME.EXT] [--overwrite] [--inplace|--out <image.imd>]
  rm      <image.imd> <CPMFILE> [--out <image.imd>]
  spdtab  <image.imd>
  spdtab  <image.imd> --drive B --class 8 [--step 0] [--target system|all] [--inplace|--out <image.imd>]

Design goals:
- Works directly on IMD (no external tools required).
- Robust for this specific IMD "style": 35 cyl, 2 heads, 9x512, MFM250k, CP/M 128B records, 2KB alloc blocks, EXM=1.
- Automatically discovers the directory track+sectors from the image.

Limitations (intentional):
- This is CP/M 2.2 directory handling only (no timestamps, no user areas >15).
- Assumes 8-bit block numbers (<=255) and 2KB blocks (inferred).
- Always rebuilds IMD with uncompressed sector records (type 0x01).

If you later want full generality (arbitrary DPB, skew tables, >256 blocks, etc.),
we can extend it — but this is meant to be a practical tool for *your* VFII images.
"""

from __future__ import annotations
import argparse
import re
import math
import itertools
import os
import uuid
from pathlib import Path
from dataclasses import dataclass
from typing import List, Optional, Tuple, Dict

# ---------------- IMD parsing ----------------

@dataclass
class IMDTrack:
    mode: int
    cyl: int
    head_raw: int
    nsec: int
    ssize_code: int
    sec_ids: List[int]
    cyl_map: Optional[List[int]]
    head_map: Optional[List[int]]
    sec_data: List[Optional[bytes]]  # None if unavailable

    @property
    def head(self) -> int:
        return self.head_raw & 0x3F

def _ssize_from_code(code: int) -> int:
    if code == 0xFF:
        raise ValueError("Per-sector sizes (0xFF) not supported in this tool.")
    return 128 << code  # 0->128, 1->256, 2->512, ...


# ---------------- VFII sector skew handling ----------------
#
# VFII CP/M disks often use a BIOS sector-translation ("skew") table even though on-disk
# sector IDs are sequential (1..N).  CP/M sees sectors in *logical* order; the physical
# sector IDs that actually contain those logical sectors are permuted.
#
# For the 9x512 MFM format used by this image, the observed skew is:
#   logical 1..9  ->  physical IDs [1,4,7,2,5,8,3,6,9]
#
# This tool builds its internal RAW representation in logical order so CP/M block/record
# math works. When writing back to IMD, we invert the mapping so the IMD keeps the original
# physical sector IDs.
# Sector translation ("skew") helpers
# RAW is built in logical-sector order; the IMD keeps the original physical sector IDs.
DEFAULT_SKEW_TABLES: Dict[int, List[int]] = {
    # VFII-style 9-sector (512B) images commonly use this 3-way interleave:
    # logical 1..9 -> physical IDs 1,4,7,2,5,8,3,6,9
    9: [1, 4, 7, 2, 5, 8, 3, 6, 9],
}

def _identity_skew_for_tracks(tracks: List["IMDTrack"]) -> List[int]:
    """Best-effort identity mapping (no translation)."""
    if not tracks:
        return []
    # Preserve the order as it appears in the IMD (physical order).
    return list(tracks[0].sec_ids)

def _parse_skew_spec(spec: str, tracks: List["IMDTrack"]) -> Optional[List[int]]:
    """
    Parse a user-provided skew spec.
      - "auto" => None (caller will auto-detect)
      - "none"/"identity" => identity mapping (no translation)
      - "default"/"vfii" => DEFAULT_SKEW_TABLES[nsec] if applicable, else identity
      - explicit list like "1,4,7,2,5,8,3,6,9" (commas or spaces)
    Returns: list of physical sector IDs in logical order, or None if "auto".
    """
    spec = (spec or "").strip().lower()
    if spec in ("", "auto"):
        return None
    if spec in ("none", "identity", "phys", "physical"):
        return _identity_skew_for_tracks(tracks)

    if not tracks:
        return []

    sample = tracks[0]
    phys_ids = list(sample.sec_ids)
    nsec = sample.nsec

    if spec in ("default", "vfii"):
        # Only apply a canned table if physical IDs are 1..N
        if sorted(phys_ids) == list(range(1, nsec + 1)):
            return list(DEFAULT_SKEW_TABLES.get(nsec, list(range(1, nsec + 1))))
        return _identity_skew_for_tracks(tracks)

    # Try explicit permutation
    # Accept comma or whitespace separated integers.
    parts = []
    for tok in spec.replace(",", " ").split():
        if tok.strip():
            parts.append(int(tok, 10))
    if len(parts) != nsec:
        raise ValueError(f"--skew expects {nsec} integers for this image (got {len(parts)}).")
    if len(set(parts)) != len(parts):
        raise ValueError("--skew must not contain duplicates.")
    # Validate against the physical IDs that actually exist in the IMD
    if set(parts) != set(phys_ids):
        raise ValueError(f"--skew must be a permutation of the track's physical sector IDs {phys_ids} (got {parts}).")
    return parts

def _candidate_skews_for_nsec(nsec: int) -> List[List[int]]:
    """Generate a small set of common skew candidates for auto-detect.

    NOTE: For VFII-style images (9 sectors), multiple candidates can score identically.
    In those tie situations we prefer the known VFII default interleave so that writes
    remain CP/M-compatible on real hardware.
    """
    base = list(range(1, nsec + 1))
    cands: List[List[int]] = []

    def add(seq: List[int]):
        if len(seq) == nsec and len(set(seq)) == nsec and all(1 <= x <= nsec for x in seq):
            if seq not in cands:
                cands.append(seq)

    # Known defaults first (tie-breaker friendly)
    if nsec in DEFAULT_SKEW_TABLES:
        add(list(DEFAULT_SKEW_TABLES[nsec]))

    # Identity (no translation)
    add(base)

    # Simple "step" interleave patterns
    for step in range(2, nsec):
        seq = []
        pos = 1
        used = set()
        for _ in range(nsec):
            while pos in used:
                pos = (pos % nsec) + 1
            seq.append(pos)
            used.add(pos)
            pos = ((pos - 1 + step) % nsec) + 1
        add(seq)

    # Factor-group interleave (covers the VFII 9-sector 3-way interleave):
    # for f|nsec: 1,1+f,... then 2,2+f,... up to f
    for f in range(2, nsec):
        if nsec % f != 0:
            continue
        seq = []
        for start in range(1, f + 1):
            for k in range(0, nsec // f):
                seq.append(start + k * f)
        add(seq)

    return cands

def _score_directory(raw: bytes, cyls: int, heads: int, bytes_per_track: int, phys_secs: int, phys_sec_size: int) -> Tuple[int, int, int, int, List[int]]:
    """
    Return best directory-likeness score over all tracks:
      (valid_active_entries, dirsec_count, dir_cyl, dir_head, logical_dir_secs)
    logical_dir_secs are 1..phys_secs indices in RAW-track order.
    """
    def looks_like_name83(e: bytes) -> bool:
        name = e[1:12]
        for ch in name:
            if ch == 0x20:  # space padding ok
                continue
            # conservative: digits + uppercase letters
            if not (0x30 <= ch <= 0x39 or 0x41 <= ch <= 0x5A):
                return False
        return True

    best = None  # (active_count, dirsec_count, c, h, sec_ids)
    for c in range(cyls):
        for h in range(heads):
            track_start = (c * heads + h) * bytes_per_track
            sec_hits: List[int] = []
            active_entries = 0
            for sid in range(1, phys_secs + 1):
                off = track_start + (sid - 1) * phys_sec_size
                sec = raw[off:off + phys_sec_size]
                if phys_sec_size == 512 and _is_dir_sector(sec):
                    sec_hits.append(sid)
                    for i in range(16):
                        e = sec[i*32:(i+1)*32]
                        if e[0] != 0xE5 and e[0] <= 0x0F and looks_like_name83(e):
                            active_entries += 1
            if len(sec_hits) >= 2 and active_entries > 0:
                score = (active_entries, len(sec_hits), c, h, sec_hits)
                if best is None or score > best:
                    best = score

    if best is None:
        return (0, 0, 0, 0, [])
    active_entries, dirsec_count, c, h, sec_hits = best
    return (active_entries, dirsec_count, c, h, sec_hits)

def autodetect_sector_skew(tracks: List["IMDTrack"], verbose: bool = False) -> List[int]:
    """
    Heuristically detect a sector-translation table (logical 1..N -> physical sector ID).

    Strategy:
      - If physical IDs are not 1..N, assume no translation (identity).
      - Otherwise, try a small set of common candidate skews, and pick the one that yields
        the most plausible CP/M directory (most valid active entries / directory sectors).
    """
    if not tracks:
        return []
    sample = tracks[0]
    nsec = sample.nsec
    phys_ids = list(sample.sec_ids)

    # If sector IDs are not 1..N, we can't safely generate candidates; keep identity.
    if sorted(phys_ids) != list(range(1, nsec + 1)):
        return _identity_skew_for_tracks(tracks)

    cyls = max(t.cyl for t in tracks) + 1
    heads = max(t.head for t in tracks) + 1
    phys_sec_size = _ssize_from_code(sample.ssize_code)
    bytes_per_track = nsec * phys_sec_size

    best = None  # (active_entries, dirsec_count, skew_list)
    for cand in _candidate_skews_for_nsec(nsec):
        raw, _ = build_raw(tracks, cand)
        active, dirsecs, _, _, dir_list = _score_directory(raw, cyls, heads, bytes_per_track, nsec, phys_sec_size)

        # Tie-breakers when directory-likeness is equal:
        # CP/M directories are typically stored in a compact, early logical region
        # (often starting at logical sector 1).
        dir_set = set(dir_list)
        dir_sorted = sorted(dir_set)
        if dir_sorted:
            rng = dir_sorted[-1] - dir_sorted[0]
            gaps = sum(1 for i in range(len(dir_sorted)-1) if (dir_sorted[i+1] - dir_sorted[i]) > 1)
            prefix = 0
            k = 1
            while k in dir_set:
                prefix += 1
                k += 1
        else:
            rng = 10**9
            gaps = 10**9
            prefix = 0

        score = (active, dirsecs, prefix, -gaps, -rng)

        if verbose:
            print(f"Auto-detect skew candidate {cand}: active_entries={active} dirsecs={dirsecs} prefix={prefix} gaps={gaps} range={rng}")

        if best is None or score > best[0]:
            best = (score, cand)

    # Require a minimum signal, otherwise fall back to default/identity.
    if best is None or best[0][0] < 2:
        return list(DEFAULT_SKEW_TABLES.get(nsec, list(range(1, nsec + 1))))
    return best[1]

def resolve_sector_skew(tracks: List["IMDTrack"], skew_spec: str, verbose: bool = False) -> List[int]:
    """Return mapping: logical sector index (1..N) -> physical sector ID."""
    parsed = _parse_skew_spec(skew_spec, tracks)
    if parsed is not None:
        return parsed
    # Auto
    return autodetect_sector_skew(tracks, verbose=verbose)

def parse_imd(path: str) -> Tuple[bytes, List[IMDTrack]]:
    blob = open(path, "rb").read()
    try:
        cut = blob.index(b"\x1A")
    except ValueError:
        raise ValueError("No 0x1A (EOF marker) found in IMD header.")
    header = blob[:cut]  # keep as-is
    p = cut + 1
    tracks: List[IMDTrack] = []

    while p < len(blob):
        mode = blob[p]; p += 1
        if p >= len(blob):
            break

        cyl = blob[p]
        head_raw = blob[p+1]
        nsec = blob[p+2]
        ssize_code = blob[p+3]
        p += 4

        has_cyl_map = bool(head_raw & 0x80)
        has_head_map = bool(head_raw & 0x40)

        sec_ids = list(blob[p:p+nsec]); p += nsec

        cyl_map = None
        head_map = None
        if has_cyl_map:
            cyl_map = list(blob[p:p+nsec]); p += nsec
        if has_head_map:
            head_map = list(blob[p:p+nsec]); p += nsec

        ssize = _ssize_from_code(ssize_code)

        sec_data: List[Optional[bytes]] = []
        for _ in range(nsec):
            rt = blob[p]; p += 1
            if rt == 0x00:
                sec_data.append(None)
                continue
            if rt in (0x01, 0x03, 0x05, 0x07):
                sec_data.append(blob[p:p+ssize]); p += ssize
                continue
            if rt in (0x02, 0x04, 0x06, 0x08):
                fill = blob[p]; p += 1
                sec_data.append(bytes([fill]) * ssize)
                continue
            raise ValueError(f"Unknown IMD sector record type: {rt:#x}")

        tracks.append(IMDTrack(mode, cyl, head_raw, nsec, ssize_code,
                               sec_ids, cyl_map, head_map, sec_data))

    return header, tracks

def build_raw(tracks: List[IMDTrack], sec_skew: List[int]) -> Tuple[bytes, Dict[Tuple[int,int], IMDTrack]]:
    """
    Build a RAW byte stream suitable for CP/M filesystem math.

    Important detail for VFII images:
      - On-disk physical sector IDs are sequential (1..N),
      - but the BIOS uses a sector-translation ("skew") table, so CP/M logical
        sector N is stored in a different physical sector ID.

    This function builds RAW in **logical sector order**:
      cyl asc, head asc, logical-sector asc
    where logical-sector asc is mapped through the skew table to a physical sector ID.

    Returns (raw_bytes, track_map).
    """
    tmap: Dict[Tuple[int,int], IMDTrack] = {(t.cyl, t.head): t for t in tracks}
    cyls = sorted({t.cyl for t in tracks})
    heads = sorted({t.head for t in tracks})
    # Assume constant geometry (true for your image)
    sample = tracks[0]
    ssize = _ssize_from_code(sample.ssize_code)
    nsec = sample.nsec
    skew = list(sec_skew)  # logical 1..N -> physical sector ID

    raw = bytearray()
    zeros = bytes([0x00]) * ssize
    for c in cyls:
        for h in heads:
            t = tmap[(c, h)]
            by_sid = {sid: (sdat if sdat is not None else zeros) for sid, sdat in zip(t.sec_ids, t.sec_data)}
            for phys_id in skew:
                raw += by_sid.get(phys_id, zeros)
    return bytes(raw), tmap

def rebuild_imd(header: bytes, tracks: List[IMDTrack], raw: bytes, sec_skew: List[int]) -> bytes:
    """
    Rebuild an IMD using original track metadata, writing all sectors as
    normal data records (type 0x01). This is simplest and compatible.

    Note: RAW is in *logical sector order* (see build_raw). When serializing back
    to IMD we invert the skew mapping so the output IMD retains the original
    physical sector IDs (typically 1..N).
    """
    tmap: Dict[Tuple[int,int], IMDTrack] = {(t.cyl, t.head): t for t in tracks}
    cyls = sorted({t.cyl for t in tracks})
    heads = sorted({t.head for t in tracks})
    skew = list(sec_skew)  # logical -> physical
    inv = {phys_id: i for i, phys_id in enumerate(skew)}  # physical -> logical index

    out = bytearray()
    out += header
    out += b"\x1A"

    p = 0
    for c in cyls:
        for h in heads:
            t = tmap[(c, h)]
            ssize = _ssize_from_code(t.ssize_code)
            nsec = t.nsec
            tbytes = raw[p:p + (ssize * nsec)]
            p += ssize * nsec

            # Split RAW track bytes into logical-sector chunks.
            logical_chunks = [tbytes[i*ssize:(i+1)*ssize] for i in range(nsec)]

            # Map physical sector ID -> bytes
            chunk_by_phys: Dict[int, bytes] = {}
            for phys_id in t.sec_ids:
                li = inv.get(phys_id, None)
                if li is None or li >= len(logical_chunks):
                    chunk_by_phys[phys_id] = bytes([0x00]) * ssize
                else:
                    chunk_by_phys[phys_id] = logical_chunks[li]

            # Track header
            out.append(t.mode)
            out.append(t.cyl)
            out.append(t.head_raw)
            out.append(t.nsec)
            out.append(t.ssize_code)
            out += bytes(t.sec_ids)
            if t.cyl_map is not None:
                out += bytes(t.cyl_map)
            if t.head_map is not None:
                out += bytes(t.head_map)

            # Sector records: emit in the same order as the IMD sector ID list
            for sid in t.sec_ids:
                out.append(0x01)
                out += chunk_by_phys[sid]

    return bytes(out)



# ---------------- VFII BIOS patching (SPDTAB) ----------------
#
# V2BIOS (and close derivatives) contain a 4-byte "Step Speed Table" (SPDTAB)
# with one entry per physical drive (A..D). Bit 7 selects disk "class":
#   0 = 8" class (also used for PC 3.5"/5.25" HD drives)
#   1 = 5" mini-floppy class
# Low 2 bits are the WD179x step-rate bits.
#
# On a V2BIOS-compatible disk, the BIOS reads the DDB from track 0 / sector 8
# using the initial density/class assumptions driven by SPDTAB. If SPDTAB says
# "5 inch" while you're using 3.5" HD media as an "8 inch" drive, you'll hit
# RNF during disk-select.
#
# This tool patches SPDTAB inside the boot/system image(s) found in the IMD by
# searching for a very stable signature: the VFII OPRTAB+ADRTAB byte table.
# In the reference sources, SPDTAB is located exactly 8 bytes before OPRTAB.
#
# Signature: OPRTAB (16 bytes) + ADRTAB (4 bytes)
_VFII_OPRTAB_ADRTAB_SIG = bytes([
    0xF0, 0x70, 0xE0, 0x60,  # C8S (side0/1, wait/no-wait)
    0xD0, 0x50, 0xC0, 0x40,  # C5S
    0xB0, 0x30, 0xA0, 0x20,  # C8D
    0x90, 0x10, 0x80, 0x00,  # C5D
    0x0E, 0x0D, 0x0B, 0x07,  # ADRTAB (A,B,C,D)
])
_VFII_SPDTAB_BACKOFFSET = 8  # bytes before OPRTAB where SPDTAB begins

@dataclass
class SpdTabHit:
    spdtab_off: int
    oprtab_off: int
    spdtab: bytes  # 4 bytes
    cyl: int
    head: int
    sec: int
    sec_off: int

def _off_to_chs(off: int, bytes_per_track: int, heads: int, sec_size: int) -> Tuple[int, int, int, int]:
    track_num = off // bytes_per_track
    in_track = off % bytes_per_track
    cyl = track_num // heads
    head = track_num % heads
    sec = in_track // sec_size + 1
    sec_off = in_track % sec_size
    return cyl, head, sec, sec_off

def find_spdtab_hits(raw: bytes, tracks: List["IMDTrack"]) -> List[SpdTabHit]:
    """
    Find SPDTAB occurrences by locating OPRTAB+ADRTAB and stepping backwards.
    Returns hits in ascending raw offset order.
    """
    if not tracks:
        return []
    sample = tracks[0]
    sec_size = _ssize_from_code(sample.ssize_code)
    heads = max(t.head for t in tracks) + 1
    bytes_per_track = sample.nsec * sec_size

    hits: List[SpdTabHit] = []
    for m in re.finditer(re.escape(_VFII_OPRTAB_ADRTAB_SIG), raw):
        oprtab_off = m.start()
        spdtab_off = oprtab_off - _VFII_SPDTAB_BACKOFFSET
        if spdtab_off < 0:
            continue
        spd = raw[spdtab_off:spdtab_off + 4]
        if len(spd) != 4:
            continue

        # Validate: only bit7 + low2 step bits should be used (bits 2..6 should be 0).
        if any((b & 0x7C) != 0 for b in spd):
            continue

        cyl, head, sec, sec_off = _off_to_chs(spdtab_off, bytes_per_track, heads, sec_size)
        hits.append(SpdTabHit(spdtab_off, oprtab_off, spd, cyl, head, sec, sec_off))

    # Deduplicate identical offsets (shouldn't happen, but be safe)
    uniq: Dict[int, SpdTabHit] = {}
    for h in hits:
        uniq.setdefault(h.spdtab_off, h)
    return [uniq[k] for k in sorted(uniq.keys())]

def _decode_spd_byte(b: int) -> str:
    klass = "5\"" if (b & 0x80) else "8\""
    step = b & 0x03
    return f"0x{b:02X} ({klass}, step={step})"

def _apply_spd_patch(val: int, klass: int, step: Optional[int]) -> int:
    """
    klass: 8 or 5
    step: optional 0..3 (WD179x step-rate bits)
    """
    if klass == 8:
        val = val & 0x7F
    elif klass == 5:
        val = val | 0x80
    else:
        raise ValueError("klass must be 8 or 5")

    if step is not None:
        if step < 0 or step > 3:
            raise ValueError("--step must be 0..3")
        val = (val & ~0x03) | (step & 0x03)
    return val & 0xFF

def _atomic_write(path: str, data: bytes) -> None:
    out_dir = os.path.dirname(os.path.abspath(path)) or "."
    tmp_path = os.path.join(out_dir, f".{os.path.basename(path)}.tmp.{uuid.uuid4().hex}")
    try:
        with open(tmp_path, "wb") as f:
            f.write(data)
            f.flush()
            os.fsync(f.fileno())
        os.replace(tmp_path, path)
    finally:
        try:
            if os.path.exists(tmp_path):
                os.remove(tmp_path)
        except OSError:
            pass

# ---------------- CP/M directory handling ----------------

DIR_ENTRY_SIZE = 32
RECORD_SIZE = 128  # CP/M "record"

@dataclass
class DiskParams:
    cyls: int
    heads: int
    phys_secs_per_track: int
    phys_sec_size: int
    bytes_per_track: int
    sec_skew: List[int]          # logical sector index (1..N) -> physical sector ID
    dir_cyl: int
    dir_head: int
    dir_phys_secs: List[int]  # logical sector numbers on that track that contain dir entries
    block_size: int           # allocation block size (bytes)
    exm: int                  # extent mask (EXM)
    block0_offset: int        # byte offset in RAW where allocation block 0 begins

def _is_dir_sector(buf512: bytes) -> bool:
    """A directory sector has 16 entries whose first byte is 0xE5 or 0..0x0F."""
    if len(buf512) != 512:
        return False
    for i in range(16):
        u = buf512[i*32]
        if not (u == 0xE5 or u <= 0x0F):
            return False
    return True

def infer_params(raw: bytes, tracks: List[IMDTrack], sec_skew: List[int]) -> DiskParams:
    # geometry from tracks
    cyls = max(t.cyl for t in tracks) + 1
    heads = max(t.head for t in tracks) + 1
    phys_secs = tracks[0].nsec
    phys_sec_size = _ssize_from_code(tracks[0].ssize_code)
    bytes_per_track = phys_secs * phys_sec_size
    sec_skew = list(sec_skew)
    # sec_skew is the mapping: logical sector index (1..N) -> physical sector ID

    # Find the track that contains the CP/M directory:
    # We scan each (cyl,head) track and look for 512-byte sectors that *could* be CP/M directory sectors.
    # A raw "all user bytes valid" heuristic can false-positive on blank/unused tracks, so we score tracks by:
    #   1) number of directory-like sectors
    #   2) number of *active* directory entries (user != 0xE5 and printable 8.3)
    def looks_like_name83(e: bytes) -> bool:
        name = e[1:12]
        if name.strip() == b"":
            return False
        for ch in name:
            if ch == 0x20:
                continue
            # CP/M allows many chars; we keep it conservative to reduce false positives
            if not (0x30 <= ch <= 0x39 or 0x41 <= ch <= 0x5A):
                return False
        return True

    best = None  # (active_count, dirsec_count, cyl, head, [sec_ids])
    for c in range(cyls):
        for h in range(heads):
            track_start = (c * heads + h) * bytes_per_track
            sec_hits: List[int] = []
            active_entries = 0
            for sid in range(1, phys_secs+1):
                off = track_start + (sid-1) * phys_sec_size
                sec = raw[off:off+phys_sec_size]
                if phys_sec_size == 512 and _is_dir_sector(sec):
                    sec_hits.append(sid)
                    # count active entries in this sector
                    for i in range(16):
                        e = sec[i*32:(i+1)*32]
                        if e[0] != 0xE5 and e[0] <= 0x0F and looks_like_name83(e):
                            active_entries += 1
            if len(sec_hits) >= 2 and active_entries > 0:
                score = (active_entries, len(sec_hits), c, h, sec_hits)
                if best is None or score > best:
                    best = score
    if best is None or best[0] < 2:
        raise ValueError("Could not locate CP/M directory sectors in this IMD.")

    active_entries, _, dir_cyl, dir_head, dir_secs = best
    # For this VFII image, directory is typically on one track and spans 2048 bytes (4 sectors).
    # We'll accept whatever we found, but require at least 2 directory sectors.
    if len(dir_secs) < 2:
        raise ValueError("Directory detection found too few sectors; image may not match expected format.")

    # Allocation parameters for this VFII CP/M image style
    # (9x512 physical, CP/M 128B records, 2KB alloc blocks, EXM=1).
    block_size = 2048
    exm = 1
    block0_offset = (dir_cyl * heads + dir_head) * bytes_per_track  # allocation block 0 begins at start of directory track

    return DiskParams(
        cyls=cyls,
        heads=heads,
        phys_secs_per_track=phys_secs,
        phys_sec_size=phys_sec_size,
        bytes_per_track=bytes_per_track,
        sec_skew=sec_skew,
        dir_cyl=dir_cyl,
        dir_head=dir_head,
        dir_phys_secs=sorted(dir_secs),
        block_size=block_size,
        exm=exm,
        block0_offset=block0_offset,
    )

def directory_entry_offsets(raw: bytes, dir_cyl: int, dir_head: int, dir_secs: List[int],
                           bytes_per_track: int, heads: int) -> List[int]:
    """
    Return byte offsets of all 32-byte entries in the detected directory sectors.
    """
    track_start = (dir_cyl * heads + dir_head) * bytes_per_track
    offsets: List[int] = []
    for sid in sorted(dir_secs):
        sec_off = track_start + (sid - 1) * 512  # this tool assumes 512-byte physical sectors
        for i in range(16):
            offsets.append(sec_off + i * DIR_ENTRY_SIZE)
    return offsets

def decode_name(e: bytes) -> str:
    fn = e[1:9].decode("ascii", "replace").rstrip()
    ex = e[9:12].decode("ascii", "replace").rstrip()
    return f"{fn}.{ex}" if ex else fn

def normalize_83(name: str) -> bytes:
    name = name.upper()
    if "." in name:
        fn, ex = name.split(".", 1)
    else:
        fn, ex = name, ""
    fn = fn[:8].ljust(8)
    ex = ex[:3].ljust(3)
    return (fn + ex).encode("ascii")

def parse_directory(raw: bytes, dp: DiskParams) -> List[Tuple[int, bytes]]:
    """
    Returns list of (entry_offset, entry_bytes) for all directory entries in directory sectors.
    """
    offsets = directory_entry_offsets(raw, dp.dir_cyl, dp.dir_head, dp.dir_phys_secs,
                                      dp.bytes_per_track, dp.heads)
    return [(o, raw[o:o+32]) for o in offsets]


def _dirent_extent_num(e: bytes) -> int:
    """Return CP/M extent number = (S2<<5) + EX (low 5 bits)."""
    ex = e[12] & 0x1F
    s2 = e[14] & 0xFF
    return (s2 << 5) | ex

def _iter_dirent_subextents(e: bytes, dp: DiskParams) -> List[Tuple[int, int, List[int]]]:
    """Decompose one directory entry into one or more logical 16K sub-extents.

    For EXM>0, CP/M can store multiple 16K sub-extents within a single directory entry,
    using different slices of the 16 block-pointer bytes. Many tools (e.g., PCPUT) will
    *pack* sub-extents this way.

    This helper is robust to both:
      - packed entries (multiple pointer slices populated)
      - split entries (only the slice corresponding to EX&EXM populated)

    Returns a list of (logical_extent_num, rc_records, blocks) sorted by logical extent.
    """
    ptrs_per_sub = 16 // (dp.exm + 1)
    extent_num = _dirent_extent_num(e)
    last_sub = (e[12] & 0x1F) & dp.exm
    start_extent = extent_num - last_sub

    out: List[Tuple[int, int, List[int]]] = []
    for sub in range(last_sub + 1):
        start = 16 + sub * ptrs_per_sub
        blks = [int(b) for b in e[start:start + ptrs_per_sub] if b != 0]
        if not blks:
            continue
        rc = 128 if sub < last_sub else int(e[15])
        out.append((start_extent + sub, rc, blks))

    out.sort(key=lambda t: t[0])
    return out

def _collect_file_extents(entries: List[bytes], dp: DiskParams) -> Dict[int, Tuple[int, List[int]]]:
    """Return map: logical_extent_num -> (rc_records, blocks).

    If duplicates exist (mixed packed + split entries), keep the version that has more
    blocks (and then larger RC).
    """
    ext_map: Dict[int, Tuple[int, List[int]]] = {}
    for e in entries:
        for ext_num, rc, blks in _iter_dirent_subextents(e, dp):
            prev = ext_map.get(ext_num)
            cand = (rc, blks)
            if prev is None:
                ext_map[ext_num] = cand
            else:
                prev_rc, prev_blks = prev
                if (len(blks), rc) > (len(prev_blks), prev_rc):
                    ext_map[ext_num] = cand
    return ext_map

def list_files(raw: bytes, dp: DiskParams) -> List[Tuple[str, int, int]]:
    """
    Return list of (NAME, size_bytes, blocks_allocated).

    Correctly handles EXM-packed directory entries (e.g., PCPUT writes EXM=1 files
    using a *single* directory entry that contains two 16K sub-extents).
    """
    dirents = parse_directory(raw, dp)

    # Collect directory entries per 8.3 name
    by_name: Dict[bytes, List[bytes]] = {}
    for _, e in dirents:
        if e[0] == 0xE5:
            continue
        if e[1:12].strip() == b"":
            continue
        key = e[1:12]
        by_name.setdefault(key, []).append(e)

    out: List[Tuple[str, int, int]] = []
    for k, ents in by_name.items():
        ext_map = _collect_file_extents(ents, dp)

        # Size is RC*128 summed across logical 16K extents.
        size = sum(rc * RECORD_SIZE for rc, _ in ext_map.values())

        # Count unique allocation blocks
        used_blocks: set[int] = set()
        for _, blks in ext_map.values():
            used_blocks.update(int(b) for b in blks)

        name = (k[:8].decode("ascii", "replace").rstrip() +
                ("." + k[8:11].decode("ascii", "replace").rstrip() if k[8:11].strip() else ""))
        out.append((name, size, len(used_blocks)))

    out.sort(key=lambda x: x[0])
    return out

def build_allocation_bitmap(raw: bytes, dp: DiskParams) -> List[bool]:
    """
    True = used, False = free for allocation blocks 0..N-1.

    IMPORTANT: CP/M EXM-packed directory entries may populate block pointers in multiple
    slices of the 16 pointer bytes. We must therefore scan *all* non-zero block pointers
    in each entry (not just one half).
    """
    total_blocks = (len(raw) - dp.block0_offset) // dp.block_size
    used = [False] * total_blocks

    # Block 0 is considered reserved/used (directory + system area starts here in this image layout).
    if total_blocks > 0:
        used[0] = True

    dirents = parse_directory(raw, dp)
    for _, e in dirents:
        if e[0] == 0xE5:
            continue
        if e[1:12].strip() == b"":
            continue
        for b in e[16:32]:
            if b == 0:
                continue
            bi = int(b)
            if bi < total_blocks:
                used[bi] = True
            else:
                raise ValueError(f"Directory references block {bi} but disk only has {total_blocks} blocks.")
    return used

def _dirent_name83_masked(e: bytes) -> bytes:
    """Return filename+ext (11 bytes) with CP/M attribute bits masked off.

    CP/M stores file attributes in the high bits of the 8+3 name bytes. Tools may set
    these bits (read-only/system/archive). For name matching, mask them.
    """
    return bytes([(b & 0x7F) for b in e[1:12]])

def find_dir_entries_for_name(dirents: List[Tuple[int, bytes]], name83: bytes, user: Optional[int] = None) -> List[Tuple[int, bytes]]:
    """Find directory entries matching an 8.3 name.

    If user is None, match any user area (0-15). Otherwise match only that user number.
    """
    out: List[Tuple[int, bytes]] = []
    for o, e in dirents:
        if e[0] == 0xE5:
            continue
        if user is not None and int(e[0]) != int(user):
            continue
        if _dirent_name83_masked(e) == name83:
            out.append((o, e))
    return out

def read_file(raw: bytes, dp: DiskParams, cpm_name: str) -> bytes:
    """Read a CP/M file from the image, handling EXM-packed extents."""
    name83 = normalize_83(cpm_name)
    dirents = parse_directory(raw, dp)
    matches = find_dir_entries_for_name(dirents, name83)
    if not matches:
        raise FileNotFoundError(f"{cpm_name} not found in directory.")

    entries = [e for _, e in matches]
    ext_map = _collect_file_extents(entries, dp)

    out = bytearray()
    for ext_num in sorted(ext_map.keys()):
        rc, blks = ext_map[ext_num]

        # Read blocks sequentially
        extent_bytes = bytearray()
        for b in blks:
            off = dp.block0_offset + int(b) * dp.block_size
            extent_bytes += raw[off:off + dp.block_size]

        out += extent_bytes[:rc * RECORD_SIZE]

    return bytes(out)

def rm_file(raw_mut: bytearray, dp: DiskParams, cpm_name: str) -> int:
    name83 = normalize_83(cpm_name)
    dirents = parse_directory(raw_mut, dp)
    matches = find_dir_entries_for_name(dirents, name83)
    if not matches:
        raise FileNotFoundError(f"{cpm_name} not found.")
    for off, _ in matches:
        raw_mut[off] = 0xE5
    return len(matches)

def put_file(raw_mut: bytearray, dp: DiskParams, local_path: str, as_name: Optional[str], overwrite: bool) -> None:
    data = open(local_path, "rb").read()
    cpm_name = as_name if as_name else (local_path.split("/")[-1].split("\\")[-1])
    name83 = normalize_83(cpm_name)

    dirents = parse_directory(raw_mut, dp)
    existing = find_dir_entries_for_name(dirents, name83)
    if existing:
        if not overwrite:
            raise ValueError(f"{cpm_name} already exists. Use --overwrite to replace.")
        # delete existing entries
        for off, _ in existing:
            raw_mut[off] = 0xE5
        # refresh directory view
        dirents = parse_directory(raw_mut, dp)

    used = build_allocation_bitmap(raw_mut, dp)
    total_blocks = len(used)

    # CP/M uses 128-byte records and 16K logical extents (128 records).
    records_total = (len(data) + RECORD_SIZE - 1) // RECORD_SIZE if len(data) else 0
    subextents_total = max(1, math.ceil(records_total / 128)) if records_total else 1

    group_size = dp.exm + 1           # e.g. EXM=1 -> 2 subextents per directory entry
    ptrs_per_sub = 16 // group_size   # e.g. 8 pointers per 16K subextent when EXM=1
    groups_needed = math.ceil(subextents_total / group_size)

    # Space needed in allocation blocks
    needed_blocks = math.ceil(len(data) / dp.block_size) if len(data) else 0
    free_blocks = [i for i in range(1, total_blocks) if not used[i]]
    if len(free_blocks) < needed_blocks:
        raise ValueError(f"Not enough free space: need {needed_blocks} blocks, have {len(free_blocks)}.")

    free_dir_offsets = [off for (off, e) in dirents if e[0] == 0xE5]
    if len(free_dir_offsets) < groups_needed:
        raise ValueError(f"Not enough free directory entries: need {groups_needed}, have {len(free_dir_offsets)}.")

    pad_byte = 0x1A  # classic CP/M padding
    data_p = 0

    def alloc_block() -> int:
        b = free_blocks.pop(0)
        used[b] = True
        return b

    for g in range(groups_needed):
        # How many 16K subextents will we store in this directory entry?
        remaining = subextents_total - g * group_size
        sub_in_group = min(group_size, remaining)
        last_sub = sub_in_group - 1
        last_extent_num = g * group_size + last_sub

        # Build directory entry
        ent = bytearray(b"\x00" * 32)
        ent[0] = 0x00          # user 0
        ent[1:12] = name83
        ent[12] = last_extent_num & 0x1F   # EX (low 5 bits)
        ent[13] = 0x00         # S1
        ent[14] = (last_extent_num >> 5) & 0xFF  # S2 (high extent bits)
        # RC is for the *last* subextent in this entry
        if records_total:
            ent[15] = min(128, max(0, records_total - last_extent_num * 128)) & 0xFF
        else:
            ent[15] = 0x00

        # Write each subextent's blocks and place pointers in the correct slice
        for sub in range(sub_in_group):
            extent_index = g * group_size + sub
            remaining_records = records_total - extent_index * 128
            rc_sub = min(128, max(0, remaining_records)) if records_total else (0 if extent_index > 0 else 0)

            need_bytes = rc_sub * RECORD_SIZE
            need_blocks = math.ceil(need_bytes / dp.block_size) if need_bytes else 0
            blocks = [alloc_block() for _ in range(need_blocks)]

            # Write blocks
            for b in blocks:
                off = dp.block0_offset + b * dp.block_size
                chunk = data[data_p:data_p + dp.block_size]
                data_p += len(chunk)
                if len(chunk) < dp.block_size:
                    chunk = chunk + bytes([pad_byte]) * (dp.block_size - len(chunk))
                raw_mut[off:off + dp.block_size] = chunk

            # Place block pointers into this subextent's slice
            start = 16 + sub * ptrs_per_sub
            for i, b in enumerate(blocks):
                ent[start + i] = b

        # Write to directory
        slot = free_dir_offsets.pop(0)
        raw_mut[slot:slot + 32] = ent

# ---------------- VFII Console I/O patching (hard-coded IN/OUT operands) ----------------
#
# In the SBC548N.* boot images from glitchwrks/versafloppy_2_cpm22, the console ports are not
# patched via a 6-byte CIOTB table; instead, the BIOS uses immediate I/O opcodes:
#
#   DB nn  -> IN  A,(nn)   ; reads USART status or data
#   D3 nn  -> OUT (nn),A   ; writes USART data
#
# The "minimal diff" patch for Bitstreamer/USART ports therefore changes only the operand
# bytes (nn). For SBC548N.IMD, there are exactly four operand bytes to patch:
#
#   0x1EB5  DB 7D   -> DB <status>
#   0x1EC2  DB 7C   -> DB <data>
#   0x1EC7  DB 7D   -> DB <status>
#   0x1ED6  D3 7C   -> D3 <data>
#
# This command defaults to patching ONLY those four offsets (safe, deterministic, matches
# your verified diffs). If you're working with a different image, you can pass --scan to
# use a heuristic search (may patch more bytes if the image contains other DB/D3 port ops).

# Known-good fixed offsets for SBC548N.IMD family:
_SBC548N_STATUS_OFFSETS = (0x1EB5, 0x1EC7)
_SBC548N_DATA_OFFSETS   = (0x1EC2, 0x1ED6)

def _parse_u8_auto(s: str) -> int:
    """Parse a byte from decimal or hex ('05', '0x05', '5')."""
    s = (s or "").strip()
    if not s:
        raise ValueError("Empty byte value.")
    if re.fullmatch(r"[0-9A-Fa-f]{1,2}", s):
        v = int(s, 16)
    else:
        v = int(s, 0)
    if not (0 <= v <= 0xFF):
        raise ValueError(f"Byte value out of range 0..255: {v}")
    return v

def _patch_fixed_sbc548n(blob_mut: bytearray, *, new_status: int, new_data: int,
                         from_status: int | None, from_data: int | None,
                         verbose: bool) -> list[tuple[int,int,int,str]]:
    """Patch exactly the 4 known SBC548N operand bytes."""
    # Verify opcodes immediately preceding the operand bytes
    # (these should be DB/DB/DB/D3 in SBC548N).
    checks = [
        (_SBC548N_STATUS_OFFSETS[0], 0xDB, "status-1"),
        (_SBC548N_DATA_OFFSETS[0],   0xDB, "data-in"),
        (_SBC548N_STATUS_OFFSETS[1], 0xDB, "status-2"),
        (_SBC548N_DATA_OFFSETS[1],   0xD3, "data-out"),
    ]
    for off, op, label in checks:
        if off <= 0 or off >= len(blob_mut):
            raise SystemExit(f"Fixed offset 0x{off:04X} out of range for this IMD.")
        if blob_mut[off-1] != op:
            raise SystemExit(
                f"Fixed-offset check failed at 0x{off:04X}: expected opcode {op:02X} "
                f"before operand ({label}), found {blob_mut[off-1]:02X}. "
                f"Use --scan for non-SBC548N images."
            )

    # Infer current ports from the fixed offsets if not specified.
    cur_status_a = blob_mut[_SBC548N_STATUS_OFFSETS[0]]
    cur_status_b = blob_mut[_SBC548N_STATUS_OFFSETS[1]]
    cur_data_a   = blob_mut[_SBC548N_DATA_OFFSETS[0]]
    cur_data_b   = blob_mut[_SBC548N_DATA_OFFSETS[1]]

    cur_status = from_status if from_status is not None else cur_status_a
    cur_data   = from_data   if from_data   is not None else cur_data_a

    # Validate consistency
    if cur_status_a != cur_status_b:
        raise SystemExit(
            f"Status operand bytes at fixed offsets do not match (0x{_SBC548N_STATUS_OFFSETS[0]:04X}="
            f"{cur_status_a:02X}, 0x{_SBC548N_STATUS_OFFSETS[1]:04X}={cur_status_b:02X}). "
            f"Use --scan."
        )
    if cur_data_a != cur_data_b:
        raise SystemExit(
            f"Data operand bytes at fixed offsets do not match (0x{_SBC548N_DATA_OFFSETS[0]:04X}="
            f"{cur_data_a:02X}, 0x{_SBC548N_DATA_OFFSETS[1]:04X}={cur_data_b:02X}). "
            f"Use --scan."
        )

    # If user supplied from_* ensure they match what's in the image.
    if from_status is not None and cur_status_a != from_status:
        raise SystemExit(
            f"--from-status says {from_status:02X} but fixed status operand is {cur_status_a:02X}. "
            f"(offset 0x{_SBC548N_STATUS_OFFSETS[0]:04X})"
        )
    if from_data is not None and cur_data_a != from_data:
        raise SystemExit(
            f"--from-data says {from_data:02X} but fixed data operand is {cur_data_a:02X}. "
            f"(offset 0x{_SBC548N_DATA_OFFSETS[0]:04X})"
        )

    if verbose:
        print(f"Fixed-offset mode (SBC548N): current status=0x{cur_status:02X}, data=0x{cur_data:02X}")
        print(f"  status offsets: {', '.join(f'0x{x:04X}' for x in _SBC548N_STATUS_OFFSETS)}")
        print(f"  data offsets:   {', '.join(f'0x{x:04X}' for x in _SBC548N_DATA_OFFSETS)}")

    patched: list[tuple[int,int,int,str]] = []

    for off in _SBC548N_STATUS_OFFSETS:
        old = blob_mut[off]
        blob_mut[off] = new_status
        patched.append((off, old, new_status, "status"))
    for off in _SBC548N_DATA_OFFSETS:
        old = blob_mut[off]
        blob_mut[off] = new_data
        patched.append((off, old, new_data, "data"))

    return patched

# --- Heuristic scan mode (optional) ---
@dataclass
class ImdSectorPayload:
    cyl: int
    head: int
    sec: int
    stype: int
    size: int
    data_off: int  # offset in IMD file where sector payload begins (only for stype 0x01/0x03/0x05/0x07)

def _iter_imd_sector_payloads(blob: bytes) -> List[ImdSectorPayload]:
    """Iterate sector payload locations in the IMD file (only full-data sectors)."""
    try:
        cut = blob.index(b"\x1A")
    except ValueError:
        raise ValueError("No 0x1A (EOF marker) found in IMD header.")
    p = cut + 1
    out: List[ImdSectorPayload] = []
    while p < len(blob):
        mode = blob[p]; p += 1
        if p >= len(blob):
            break
        cyl = blob[p]
        head_raw = blob[p+1]
        nsec = blob[p+2]
        ssize_code = blob[p+3]
        p += 4
        has_cyl_map = bool(head_raw & 0x80)
        has_head_map = bool(head_raw & 0x40)
        sec_ids = list(blob[p:p+nsec]); p += nsec
        if has_cyl_map:
            cyl_map = list(blob[p:p+nsec]); p += nsec
        else:
            cyl_map = [cyl] * nsec
        if has_head_map:
            head_map = list(blob[p:p+nsec]); p += nsec
        else:
            head_map = [head_raw & 0x3F] * nsec
        ssize = _ssize_from_code(ssize_code)
        for i in range(nsec):
            stype = blob[p]; p += 1
            if stype == 0x00:
                continue
            if stype in (0x01, 0x03, 0x05, 0x07):
                out.append(ImdSectorPayload(cyl_map[i], head_map[i], sec_ids[i], stype, ssize, p))
                p += ssize
            elif stype in (0x02, 0x04, 0x06, 0x08):
                p += 1
            else:
                raise ValueError(f"Unknown IMD sector record type: {stype:#x}")
    return out

def _scan_and_patch(blob_mut: bytearray, *, new_status: int, new_data: int,
                    from_status: int | None, from_data: int | None,
                    verbose: bool) -> list[tuple[int,int,int,str]]:
    """Heuristic scan: may match more than intended on some images."""
    payloads = _iter_imd_sector_payloads(bytes(blob_mut))

    # Find likely operand positions for status (DB xx E6 01/02), IN (DB xx with nearby RET), OUT (D3 xx)
    status_ops: list[int] = []
    in_ops: list[int] = []
    out_ops: list[int] = []

    for sec in payloads:
        buf = blob_mut[sec.data_off:sec.data_off + sec.size]

        # Status patterns: DB xx E6 01 / DB xx E6 02
        for mask in (0x01, 0x02):
            start = 0
            while True:
                idx = buf.find(bytes([0xDB]), start)
                if idx == -1 or idx + 4 > len(buf):
                    break
                if buf[idx+2] == 0xE6 and buf[idx+3] == mask:
                    status_ops.append(sec.data_off + idx + 1)
                start = idx + 1

        # OUT patterns: D3 xx
        start = 0
        while True:
            idx = buf.find(bytes([0xD3]), start)
            if idx == -1 or idx + 2 > len(buf):
                break
            out_ops.append(sec.data_off + idx + 1)
            start = idx + 1

        # IN patterns: DB xx with RET nearby (C9)
        start = 0
        while True:
            idx = buf.find(bytes([0xDB]), start)
            if idx == -1 or idx + 2 > len(buf):
                break
            window = buf[idx+2:idx+12]
            if 0xC9 in window:
                in_ops.append(sec.data_off + idx + 1)
            start = idx + 1

    # Dedup
    def _dedup(xs: list[int]) -> list[int]:
        out = []
        seen = set()
        for x in xs:
            if x not in seen:
                out.append(x)
                seen.add(x)
        return out
    status_ops = _dedup(status_ops)
    in_ops = _dedup(in_ops)
    out_ops = _dedup(out_ops)

    # Infer current ports if not supplied: use modes
    def _mode(vals: list[int]) -> int:
        return max(set(vals), key=vals.count)

    if not status_ops:
        raise SystemExit("Scan mode: no status patterns found.")
    cur_status = from_status if from_status is not None else _mode([blob_mut[o] for o in status_ops])

    data_candidates = []
    data_candidates += [blob_mut[o] for o in out_ops if blob_mut[o] != cur_status]
    data_candidates += [blob_mut[o] for o in in_ops if blob_mut[o] != cur_status]
    if not data_candidates:
        raise SystemExit("Scan mode: could not infer data port.")
    cur_data = from_data if from_data is not None else _mode(data_candidates)

    if verbose:
        print(f"Scan mode: inferred status=0x{cur_status:02X}, data=0x{cur_data:02X}")
        print(f"  status hits: {len(status_ops)}  in hits: {len(in_ops)}  out hits: {len(out_ops)}")

    patched: list[tuple[int,int,int,str]] = []
    for o in status_ops:
        if blob_mut[o] == cur_status:
            old = blob_mut[o]
            blob_mut[o] = new_status
            patched.append((o, old, new_status, "status"))
    for o in in_ops:
        if blob_mut[o] == cur_data:
            old = blob_mut[o]
            blob_mut[o] = new_data
            patched.append((o, old, new_data, "data-in"))
    for o in out_ops:
        if blob_mut[o] == cur_data:
            old = blob_mut[o]
            blob_mut[o] = new_data
            patched.append((o, old, new_data, "data-out"))

    # Dedup
    uniq = {p[0]: p for p in patched}
    return [uniq[k] for k in sorted(uniq.keys())]

def cmd_consolepatch(args: argparse.Namespace) -> None:
    """Patch console I/O ports in an IMD."""
    blob_mut = bytearray(Path(args.imd).read_bytes())
    new_status = _parse_u8_auto(args.status)
    new_data = _parse_u8_auto(args.data)
    from_status = _parse_u8_auto(args.from_status) if args.from_status is not None else None
    from_data = _parse_u8_auto(args.from_data) if args.from_data is not None else None

    if args.scan:
        patched = _scan_and_patch(blob_mut, new_status=new_status, new_data=new_data,
                                  from_status=from_status, from_data=from_data,
                                  verbose=args.verbose)
    else:
        patched = _patch_fixed_sbc548n(blob_mut, new_status=new_status, new_data=new_data,
                                       from_status=from_status, from_data=from_data,
                                       verbose=args.verbose)

    if not patched:
        raise SystemExit("No bytes were patched (already set, or patterns not found).")

    if args.inplace and args.out:
        raise SystemExit("Error: --inplace and --out are mutually exclusive.")

    if args.inplace:
        _atomic_write(args.imd, bytes(blob_mut))
        out_path = args.imd
    else:
        out_path = args.out if args.out else str(Path(args.imd).with_name(
            f"{Path(args.imd).stem}_serial_{new_data:02X}_{new_status:02X}{Path(args.imd).suffix}"
        ))
        Path(out_path).write_bytes(blob_mut)

    print(f"Console patch: status -> 0x{new_status:02X}, data -> 0x{new_data:02X}")
    print(f"Patched {len(patched)} byte(s). Wrote {out_path}")
    if args.verbose:
        for o, old, new, kind in patched:
            print(f"  off=0x{o:05X}: {old:02X}->{new:02X} ({kind})")


# ---------------- CLI ----------------


def cmd_spdtab(args: argparse.Namespace) -> None:
    """
    Show and/or patch V2BIOS SPDTAB entries inside boot/system image(s) embedded in the IMD.

    By default (no --drive), this prints all SPDTAB tables found.

    Example:
      # Show:
      python3 imd_cpmtools_vfii.py spdtab disk.imd

      # Patch drive B to "8 inch class", keep existing step bits:
      python3 imd_cpmtools_vfii.py spdtab disk.imd --drive B --class 8 --target system --out disk_spdtabB8.imd

      # Patch drive B to "8 inch class", force fastest step bits (0):
      python3 imd_cpmtools_vfii.py spdtab disk.imd --drive B --class 8 --step 0 --inplace
    """
    header, tracks = parse_imd(args.imd)
    sec_skew = resolve_sector_skew(tracks, args.skew, verbose=getattr(args, 'skew_verbose', False))
    raw, _ = build_raw(tracks, sec_skew)

    hits = find_spdtab_hits(raw, tracks)
    if not hits:
        raise SystemExit("No SPDTAB/OPRTAB signature found in this IMD. (Not a V2BIOS-style system image?)")

    # Try to infer directory track so we can label early-track hits as "system".
    dir_cyl = None
    try:
        dp = infer_params(raw, tracks, sec_skew)
        dir_cyl = dp.dir_cyl
    except Exception:
        dir_cyl = None

    def tag_for(h: SpdTabHit) -> str:
        if dir_cyl is None:
            return ""
        return " [system]" if h.cyl < dir_cyl else " [data]"

    print(f"Found {len(hits)} SPDTAB table(s) (signature match).")
    for i, h in enumerate(hits, 1):
        spd = list(h.spdtab)
        print(f"#{i}: raw_off=0x{h.spdtab_off:05X}  C={h.cyl} H={h.head}  sec={h.sec} +{h.sec_off}{tag_for(h)}")
        print(f"    A={_decode_spd_byte(spd[0])}  B={_decode_spd_byte(spd[1])}  C={_decode_spd_byte(spd[2])}  D={_decode_spd_byte(spd[3])}")

    # Show-only mode
    if not args.drive:
        return

    if args.klass is None:
        raise SystemExit("Error: --class is required when using --drive (use 8 or 5).")

    drive_idx = {"A": 0, "B": 1, "C": 2, "D": 3}[args.drive.upper()]
    klass = int(args.klass)

    # Select which hits to patch.
    selected: List[SpdTabHit] = []
    if args.target == "first":
        selected = hits[:1]
    elif args.target == "all":
        selected = hits
    elif args.target == "system":
        if dir_cyl is None:
            selected = hits[:1]  # best-effort: patch the earliest copy
        else:
            selected = [h for h in hits if h.cyl < dir_cyl]
    elif args.target == "files":
        if dir_cyl is None:
            selected = hits[1:] if len(hits) > 1 else []
        else:
            selected = [h for h in hits if h.cyl >= dir_cyl]
    else:
        raise SystemExit(f"Unknown --target {args.target}")

    if not selected:
        raise SystemExit(f"No SPDTAB hits matched --target {args.target}.")

    raw_mut = bytearray(raw)
    changes = []
    for h in selected:
        off = h.spdtab_off + drive_idx
        old = raw_mut[off]
        new = _apply_spd_patch(old, klass=klass, step=args.step)
        raw_mut[off] = new
        changes.append((h, old, new))

    print("")
    for h, old, new in changes:
        print(f"Patched at raw_off=0x{h.spdtab_off + drive_idx:05X}  C={h.cyl} H={h.head}  "
              f"{args.drive.upper()}: {_decode_spd_byte(old)} -> {_decode_spd_byte(new)}{tag_for(h)}")

    out_imd = rebuild_imd(header, tracks, bytes(raw_mut), sec_skew)

    if getattr(args, "inplace", False) and args.out:
        raise SystemExit("Error: --inplace and --out are mutually exclusive.")

    if getattr(args, "inplace", False):
        _atomic_write(args.imd, out_imd)
        out_path = args.imd
        print(f"\nUpdated in-place: {out_path}")
    else:
        if args.out:
            out_path = args.out
        else:
            p = Path(args.imd)
            out_path = str(p.with_name(f"{p.stem}_spdtab_patched{p.suffix}"))
        with open(out_path, "wb") as f:
            f.write(out_imd)
        print(f"\nWrote {out_path}")

    # Quick re-scan summary
    header2, tracks2 = parse_imd(out_path)
    sec_skew2 = resolve_sector_skew(tracks2, args.skew, verbose=getattr(args, 'skew_verbose', False))
    raw2, _ = build_raw(tracks2, sec_skew2)
    hits2 = find_spdtab_hits(raw2, tracks2)
    if hits2:
        print("\nPost-patch SPDTAB (first hit shown):")
        spd = list(hits2[0].spdtab)
        print(f"    A={_decode_spd_byte(spd[0])}  B={_decode_spd_byte(spd[1])}  C={_decode_spd_byte(spd[2])}  D={_decode_spd_byte(spd[3])}")


def cmd_ls(args: argparse.Namespace) -> None:
    header, tracks = parse_imd(args.imd)
    sec_skew = resolve_sector_skew(tracks, args.skew, verbose=getattr(args, 'skew_verbose', False))
    raw, _ = build_raw(tracks, sec_skew)
    dp = infer_params(raw, tracks, sec_skew)
    files = list_files(raw, dp)
    for name, size, blocks in files:
        print(f"{name:12} {size:8} bytes  ({blocks} blocks)")

def cmd_get(args: argparse.Namespace) -> None:
    header, tracks = parse_imd(args.imd)
    sec_skew = resolve_sector_skew(tracks, args.skew, verbose=getattr(args, 'skew_verbose', False))
    raw, _ = build_raw(tracks, sec_skew)
    dp = infer_params(raw, tracks, sec_skew)
    data = read_file(raw, dp, args.cpmfile)
    out_path = args.out if args.out else args.cpmfile
    with open(out_path, "wb") as f:
        f.write(data)
    print(f"Wrote {len(data)} bytes to {out_path}")

def cmd_rm(args: argparse.Namespace) -> None:
    header, tracks = parse_imd(args.imd)
    sec_skew = resolve_sector_skew(tracks, args.skew, verbose=getattr(args, 'skew_verbose', False))
    raw, _ = build_raw(tracks, sec_skew)
    dp = infer_params(raw, tracks, sec_skew)
    raw_mut = bytearray(raw)
    n = rm_file(raw_mut, dp, args.cpmfile)
    out_imd = rebuild_imd(header, tracks, bytes(raw_mut), sec_skew)
    out_path = args.out if args.out else args.imd.replace(".IMD", "_patched.IMD")
    with open(out_path, "wb") as f:
        f.write(out_imd)
    print(f"Removed {n} directory entries for {args.cpmfile}. Wrote {out_path}")

def cmd_put(args: argparse.Namespace) -> None:
    header, tracks = parse_imd(args.imd)
    sec_skew = resolve_sector_skew(tracks, args.skew, verbose=getattr(args, 'skew_verbose', False))
    raw, _ = build_raw(tracks, sec_skew)
    dp = infer_params(raw, tracks, sec_skew)
    raw_mut = bytearray(raw)

    cpm_name = args.as_name if args.as_name else os.path.basename(args.local)
    put_file(raw_mut, dp, args.local, cpm_name, args.overwrite)

    out_imd = rebuild_imd(header, tracks, bytes(raw_mut), sec_skew)

    # Output behavior:
    #   - default: write a new IMD (like before)
    #   - --inplace: update the input IMD path (written via temp + atomic replace for safety)
    if getattr(args, "inplace", False) and args.out:
        raise SystemExit("Error: --inplace and --out are mutually exclusive.")

    if getattr(args, "inplace", False):
        out_path = args.imd
        out_dir = os.path.dirname(os.path.abspath(out_path)) or "."
        tmp_path = os.path.join(out_dir, f".{os.path.basename(out_path)}.tmp.{uuid.uuid4().hex}")
        try:
            with open(tmp_path, "wb") as f:
                f.write(out_imd)
                f.flush()
                os.fsync(f.fileno())
            os.replace(tmp_path, out_path)
        finally:
            # If something went wrong before replace, best-effort cleanup.
            try:
                if os.path.exists(tmp_path):
                    os.remove(tmp_path)
            except OSError:
                pass
        print(f"Inserted {args.local} as {cpm_name}. Updated in-place: {out_path}")
    else:
        out_path = args.out if args.out else args.imd.replace(".IMD", "_patched.IMD")
        with open(out_path, "wb") as f:
            f.write(out_imd)
        print(f"Inserted {args.local} as {cpm_name}. Wrote {out_path}")

    if args.verify:
        # Re-open the newly written image and verify that the stored file matches the input.
        header2, tracks2 = parse_imd(out_path)
        sec_skew2 = resolve_sector_skew(tracks2, args.skew, verbose=getattr(args, 'skew_verbose', False))
        raw2, _ = build_raw(tracks2, sec_skew2)
        dp2 = infer_params(raw2, tracks2, sec_skew2)

        stored = read_file(raw2, dp2, cpm_name)
        inp = Path(args.local).read_bytes()
        ok_prefix = stored[:len(inp)] == inp
        ok_pad = all(b in (0x00, 0x1A) for b in stored[len(inp):])

        # Locate directory entry/entries for quick visibility.
        dirents = parse_directory(raw2, dp2)
        n83 = normalize_83(cpm_name)
        matches = [(off, e) for off, e in dirents if e[0] != 0xE5 and e[1:12] == n83]
        if matches:
            # Show each directory entry and the subextents we detect in it.
            for off, e in sorted(matches, key=lambda t: (_dirent_extent_num(t[1]), t[0])):
                subs = _iter_dirent_subextents(e, dp2)
                # Flatten blocks for display
                all_blocks = [b for _, _, bl in subs for b in bl]
                print(f"Verify: dirent_off=0x{off:04X} EX={e[12]} S1={e[13]} S2={e[14]} RC={e[15]} blocks={all_blocks} subextents={[(en, rc) for en, rc, _ in subs]}")

        if ok_prefix and ok_pad:
            print("Verify: OK (stored bytes match input; remaining record padding is 0x00/0x1A).")
        else:
            print("Verify: FAILED")
            print(f"  stored_len={len(stored)} input_len={len(inp)}")
            if not ok_prefix:
                # Find first mismatch index (within input length).
                mm = next((i for i in range(min(len(stored), len(inp))) if stored[i] != inp[i]), None)
                print(f"  first_mismatch_at={mm}")
            if not ok_pad:
                print("  unexpected padding bytes found after input length (expected 0x00 or 0x1A).")

def cmd_info(args):
    header, tracks = parse_imd(args.imd)
    sec_skew = resolve_sector_skew(tracks, args.skew, verbose=getattr(args, 'skew_verbose', False))
    raw, _ = build_raw(tracks, sec_skew)
    dp = infer_params(raw, tracks, sec_skew)

    header_text = header.decode("ascii", errors="replace")
    first_line = header_text.splitlines()[0] if header_text.splitlines() else "(no header text)"

    cyls = dp.cyls
    heads = dp.heads
    secs = dp.phys_secs_per_track
    ssize = dp.phys_sec_size
    expected_bytes = cyls * heads * secs * ssize

    modes = {}
    for t in tracks:
        modes[t.mode] = modes.get(t.mode, 0) + 1

    def mode_desc(m: int) -> str:
        # ImageDisk modes:
        # 0=500k FM, 1=300k FM, 2=250k FM, 3=500k MFM, 4=300k MFM, 5=250k MFM
        return {
            0: "FM 500kbps",
            1: "FM 300kbps",
            2: "FM 250kbps",
            3: "MFM 500kbps",
            4: "MFM 300kbps",
            5: "MFM 250kbps",
        }.get(m, "unknown")

    files = list_files(raw, dp)
    used = build_allocation_bitmap(raw, dp)
    total_blocks = len(used)
    free_blocks = sum(1 for u in used if not u)
    used_blocks = total_blocks - free_blocks

    print(f"IMD: {args.imd}")
    print(f"Header: {first_line}")
    # Report the chosen sector skew mapping (logical sector -> physical sector ID)
    spec = (getattr(args, "skew", "auto") or "auto").strip().lower()
    if spec in ("", "auto"):
        skew_note = "auto"
    elif spec in ("none", "identity", "phys", "physical"):
        skew_note = "identity"
    elif spec in ("default", "vfii"):
        skew_note = "vfii/default"
    else:
        skew_note = "explicit"
    skew_str = ", ".join(str(x) for x in sec_skew)
    print(f"Selected skew (logical->physical): {skew_str} ({skew_note})")
    print("")
    print("Geometry (physical, from IMD):")
    print(f"  Cylinders : {cyls}")
    print(f"  Heads     : {heads}")
    print(f"  Sectors/tr: {secs}")
    print(f"  Sec size  : {ssize} bytes")
    print("  Track modes: " + ", ".join(f"{k}:{mode_desc(k)} x{v}" for k, v in sorted(modes.items())))
    print(f"  Raw size  : {len(raw)} bytes (expected {expected_bytes} bytes)")
    if len(raw) != expected_bytes:
        print("  NOTE: Raw size differs from simple cyl*head*sec*size; this can happen if IMD has missing tracks.")

    print("")
    print("CP/M layout (inferred for this image):")
    print(f"  Directory track: C={dp.dir_cyl} H={dp.dir_head}")
    dir_log = sorted(dp.dir_phys_secs)
    dir_phys = [dp.sec_skew[s-1] for s in dir_log if 1 <= s <= len(dp.sec_skew)]
    print(f"  Directory sectors (logical): {', '.join(str(s) for s in dir_log)}")
    print(f"  Directory sectors (physical IDs): {', '.join(str(s) for s in dir_phys)}")
    print(f"  Allocation block size: {dp.block_size} bytes")
    print(f"  EXM: {dp.exm}  (block pointers per extent: {16 // (dp.exm + 1)})")
    print(f"  Block0 offset into RAW: {dp.block0_offset} bytes")

    print("")
    print("Disk usage:")
    print(f"  Total blocks: {total_blocks}")
    print(f"  Used blocks : {used_blocks}")
    print(f"  Free blocks : {free_blocks}")
    print(f"  Free bytes  : {free_blocks * dp.block_size}")
    print(f"  File count  : {len(files)}")

    if args.long:
        print("")
        print("Files:")
        for name, size, blocks in files:
            print(f"  {name:<12} {size:>7} bytes  blocks={blocks}")


def main():
    ap = argparse.ArgumentParser(description="CP/M file tools for VFII-style IMD images.")
    sub = ap.add_subparsers(dest="cmd", required=True)
    common = argparse.ArgumentParser(add_help=False)
    common.add_argument(
        "--skew",
        default="auto",
        help="Sector skew (logical->physical). Use 'auto' (default), 'identity', 'default'/'vfii', or an explicit list like '1,4,7,2,5,8,3,6,9'.",
    )
    common.add_argument(
        "--skew-verbose",
        dest="skew_verbose",
        action="store_true",
        help="Print auto-detect scoring for skew candidates.",
    )


    p_ls = sub.add_parser("ls", help="List CP/M directory", parents=[common])
    p_ls.add_argument("imd")
    p_ls.set_defaults(func=cmd_ls)

    p_info = sub.add_parser("info", help="Show geometry + inferred CP/M parameters + free space", parents=[common])
    p_info.add_argument("imd")
    p_info.add_argument("--long", action="store_true", help="Also list files with sizes")
    p_info.set_defaults(func=cmd_info)


    p_spd = sub.add_parser(
        "spdtab",
        help="Show or patch V2BIOS SPDTAB (drive class/step bits) inside boot/system image(s)",
        parents=[common],
    )
    p_spd.add_argument("imd")
    p_spd.add_argument("--drive", choices=["A", "B", "C", "D"], default=None, help="Drive letter to modify. If omitted, just show tables.")
    p_spd.add_argument("--class", dest="klass", choices=["8", "5"], default=None, help="Disk class to set for that drive in SPDTAB: 8 or 5.")
    p_spd.add_argument("--step", type=int, choices=[0, 1, 2, 3], default=None, help="Optional WD179x step-rate bits (0..3).")
    p_spd.add_argument("--target", choices=["all", "system", "files", "first"], default="all", help="Which SPDTAB copies to patch if multiple are found.")
    p_spd.add_argument("--out", default=None, help="Output IMD path (default: *_spdtab_patched.IMD)")
    p_spd.add_argument("--inplace", action="store_true", help="Update the input IMD file in-place (uses temp + atomic replace).")
    p_spd.set_defaults(func=cmd_spdtab)
    p_con = sub.add_parser("console", help="Patch console I/O ports in a VFII IMD (SBC548N fixed offsets by default)", parents=[common])
    p_con.add_argument("imd")
    p_con.add_argument("--status", required=True, help="New status port (e.g. 05 or 0x05)")
    p_con.add_argument("--data", required=True, help="New data port (e.g. 04 or 0x04)")
    p_con.add_argument("--from-status", dest="from_status", default=None, help="Expected current status port (optional safety check)")
    p_con.add_argument("--from-data", dest="from_data", default=None, help="Expected current data port (optional safety check)")
    p_con.add_argument("--scan", action="store_true", help="Use heuristic scan instead of SBC548N fixed offsets (may patch more bytes).")
    p_con.add_argument("--out", default=None, help="Output IMD path (default: *_serial_DD_SS.IMD)")
    p_con.add_argument("--inplace", action="store_true", help="Update the input IMD file in-place (uses temp + atomic replace).")
    p_con.add_argument("--verbose", action="store_true", help="Print patch details.")
    p_con.set_defaults(func=cmd_consolepatch)


    p_get = sub.add_parser("get", help="Extract a CP/M file", parents=[common])
    p_get.add_argument("imd")
    p_get.add_argument("cpmfile", help="CP/M 8.3 name, e.g. V2FORMAT.COM")
    p_get.add_argument("--out", help="Output path (default: CP/M name)", default=None)
    p_get.set_defaults(func=cmd_get)

    p_put = sub.add_parser("put", help="Insert a local file into the CP/M directory", parents=[common])
    p_put.add_argument("imd")
    p_put.add_argument("local", help="Local path to file")
    p_put.add_argument("--as", dest="as_name", help="CP/M 8.3 name to store as", default=None)
    p_put.add_argument("--overwrite", action="store_true", help="Replace if file already exists")
    p_put.add_argument("--out", help="Output IMD path (default: *_patched.IMD)", default=None)
    p_put.add_argument("--inplace", action="store_true", help="Update the input IMD file instead of writing a new one (uses temp + atomic replace).")
    p_put.add_argument("--verify", action="store_true", help="After writing, re-open output IMD and verify stored bytes match the input file")
    p_put.set_defaults(func=cmd_put)

    p_rm = sub.add_parser("rm", help="Delete a CP/M file (marks dir entries as 0xE5)", parents=[common])
    p_rm.add_argument("imd")
    p_rm.add_argument("cpmfile")
    p_rm.add_argument("--out", help="Output IMD path (default: *_patched.IMD)", default=None)
    p_rm.set_defaults(func=cmd_rm)

    args = ap.parse_args()
    args.func(args)

if __name__ == "__main__":
    main()
