#!/usr/bin/env python3
# build.py - build the complete annotated disassembly of Modem Wars (C64) from the decrypted disk.
#
#   python3 tools/build.py            -> writes disassembly/ (sources, includes, verify script)
#
# Annotation files in annotations/*.json are merged in (labels, comments, data types, zero page names).
import os, sys, json, glob
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from m6502 import trace, decode, ZP, ZPX, ZPY, ABS, ABX, ABY, IND, IZX, IZY, REL
from render import RendererT, writeIncludes
from hwSymbols import c64Symbols, DRIVE, KERNAL
from discover import instructionStarts, vectorEntries, patchedCallTables, wordsFrom, splitTableEntries

ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
EXT  = os.path.join(ROOT, "extracted")
OUT  = os.path.join(ROOT, "disassembly")
ANN  = os.path.join(ROOT, "annotations")

def sec(t, s):
    return open(f"{EXT}/decrypted/t{t:02d}s{s:02d}.bin", "rb").read()

def rawSec(t, s):
    return open(f"{EXT}/sectors/t{t:02d}s{s:02d}.bin", "rb").read()

def runOf(t, s0, n):
    return b"".join(sec(t, s0 + i) for i in range(n))

bootImage = bytearray(open(f"{EXT}/bootImage.bin", "rb").read())
loadPrg   = open(f"{EXT}/files/load.prg", "rb").read()[2:]
eaPrg     = open(f"{EXT}/files/ea.prg", "rb").read()

# ---------------------------------------------------------------------------------------------
# Overlay definitions: name -> (loadAddr, bytes, provenance text)
# ---------------------------------------------------------------------------------------------
OVL = {
 "C000_T30":   (0xC000, runOf(30,0,16),                  "track 30 sectors 0-15"),
 "6F00_A":     (0x6F00, runOf(31,0,17) + runOf(32,0,8),  "track 31 sectors 0-16 + track 32 sectors 0-7"),
 "6F00_B":     (0x6F00, runOf(33,0,17) + runOf(32,8,8),  "track 33 sectors 0-16 + track 32 sectors 8-15"),
 "8800_T29":   (0x8800, runOf(29,9,2),                   "track 29 sectors 9-10"),
 "E000_T35":   (0xE000, runOf(35,0,16),                  "track 35 sectors 0-15"),
 "E000_T34":   (0xE000, sec(18,7) + runOf(34,1,15),      "track 18 sector 7 + track 34 sectors 1-15"),
 "EC00_T29":   (0xEC00, runOf(29,0,4),                   "track 29 sectors 0-3"),
 "EE00_T29":   (0xEE00, runOf(29,4,2),                   "track 29 sectors 4-5"),
 "F000_T18":   (0xF000, runOf(18,8,7),                   "track 18 sectors 8-14"),
 "0200_T18":   (0x0200, runOf(18,15,2),                  "track 18 sectors 15-16"),
 "0200_T29":   (0x0200, runOf(29,6,3),                   "track 29 sectors 6-8"),
}

def makeImage(*names):
    img = bytearray(bootImage)
    for n in names:
        addr, data, _ = OVL[n]
        img[addr:addr+len(data)] = data
    return img

# runtime image variants used for tracing (each overlay traced inside a compatible image)
IMAGES = {
 "B":  makeImage("C000_T30", "6F00_B", "8800_T29", "E000_T35", "F000_T18", "0200_T18"),
 "A":  makeImage("C000_T30", "6F00_A", "8800_T29", "E000_T34", "F000_T18", "0200_T18"),
 "A2": makeImage("C000_T30", "6F00_A", "8800_T29", "E000_T34", "EC00_T29", "F000_T18", "0200_T18"),
 "B3": makeImage("C000_T30", "6F00_B", "8800_T29", "E000_T35", "EE00_T29", "F000_T18", "0200_T18"),
 "D":  makeImage("C000_T30", "6F00_A", "8800_T29", "E000_T34", "EC00_T29", "F000_T18", "0200_T29"),
}
RUNTIME_VALID = [(0x0200,0x8A00),(0x8C00,0x9000),(0x9300,0xD000),(0xE000,0xF700),(0xFBB8,0xFFD2)]
BOOT_VALID    = [(0x02A8,0x030C),(0xC000,0xC400)]

def validFor(ranges):
    return lambda a: any(lo <= a < hi for lo, hi in ranges)

# ---------------------------------------------------------------------------------------------
# Units (output files)
# ---------------------------------------------------------------------------------------------
UNITS = [
 # name, image key, start, end, title, provenance lines, extra entry points
 ("boot/eaBootFile",    "BOOTEA", 0x02A8, 0x030C, "EA - 1-block BASIC-vector hijack loader (directory file \"ea\", loads at $02A8)",
    ["Loaded with LOAD\"EA\",8,1 (or LOAD\"*\",8,1). The file overwrites the BASIC vectors at $0300-$030B so that",
     "control reaches $02B8 as soon as BASIC returns to its main loop after the LOAD completes.",
     "$02B8 then KERNAL-LOADs the file \"load\" ($9800-$C3FF) and calls the fast loader at $C000."], [0x02B8]),
 ("boot/fastLoaderC000", "BOOT",  0xC000, 0xC400, "LOAD ($C000-$C3FF) - Electronic Arts fast loader / bootstrap, C64 side",
    ["Part of directory file \"load\" (PRG, $9800-$C3FF). Entry point $C145 is called once from $02B8.",
     "Installs drive code via a DOS \"B-E\" (block execute) of track 1 sector 17, then pulls the game from",
     "the disk through a 2-bit serial protocol on CIA2 port A ($DD00); the drive decrypts each sector before",
     "sending it, so every sector except those on track 18 is stored XOR-encrypted on disk.",
     "$C004/$C034 form the public API the boot sequence uses; the game keeps a private copy at $0804/$085C",
     "because this page is later overwritten by the text-engine overlay."],
    [0xC145, 0xC004, 0xC034, 0xC229, 0xC3C4]),
 ("boot/titleColorRam9800", "LOADPRG", 0x9800, 0x9C00, "LOAD ($9800-$9BFF) - title picture colour RAM (1000 bytes, copied to $D800 by $C23F)", ["Low nibble = colour 3 of each 4x8 multicolour cell. Bytes $9BE8-$9BFF are padding."], []),
 ("boot/titleScreenRam9C00", "LOADPRG", 0x9C00, 0xA000, "LOAD ($9C00-$9FFF) - title picture screen RAM (copied to $8C00 by $C23F; VIC bank 2, $D018=$38)", ["High nibble = colour 1, low nibble = colour 2 of each multicolour cell."], []),
 ("boot/titleBitmapA000", "LOADPRG", 0xA000, 0xC000, "LOAD ($A000-$BF3F) - title picture multicolour bitmap (8000 bytes, VIC bank 2, $D011=$3B, $D016=$D8)", ["Displayed by the boot loader while the game loads. $BF40-$BFFF is padding."], []),
 ("game/tileRules0200",  "B",      0x0200, 0x0400, "$0200-$03FF - terrain tile rule table and start-up code (track 18 sectors 15-16)",
    ["Loaded by $10BB. The 256-byte table at $0200 maps a neighbourhood bit mask (bit i = neighbour i, going",
     "NW, N, NE, E, SE, S, SW, W) to the contour tile the map generator must draw; $0300 holds start-up code",
     "that the main program calls at $0A87. The area is reused as a sprite-shape decompression buffer once",
     "the battle starts."], []),
 ("game/trainerPlaybook0200", "D", 0x0200, 0x0500, "$0200-$04FF - solo trainer playbook (track 29 sectors 6-8), loaded by the trainer setup overlay at $ECAA",
    ["Five canned opening plays; the trainer setup code copies a 150-byte slice into the AI's group and",
     "destination tables. It overwrites the tile rule table and the first page of the scratch area."], []),
 ("game/scratchAndMessages0400", "B", 0x0400, 0x0800, "$0400-$07FF - manual-check answer table, status message strings and shared scratch area",
    ["Copied from $6F00 (track 27 sectors 13-16) by the boot loader. $0401 holds the 32 x 3-byte answers for",
     "the manual look-up check, $0461 the memory checksum routine the boot loader returns into, $0500-$0567",
     "the status message pointer tables and $056C.. the status message strings. Once the battle starts the",
     "three 50-byte arrays at $0400/$0432/$0464 are reused as unit scan lists and AI threat maps."], []),
 ("game/mainProgram0800", "B",    0x0800, 0x6F00, "$0800-$6EFF - main game program (tracks 22-27, sectors 0-17 each)",
    ["Entry point $0800 (JMP $0A6C). Contains the loader API, start-up sequence, raster interrupt, input",
     "handling, the battlefield and console screens, the unit simulation and the command interpreter."], [0x0800]),
 ("game/mapGenerator6F00", "A",   0x6F00, 0x8800, "$6F00-$87FF - random battlefield generator and setup phase (track 31 sectors 0-16 + track 32 sectors 0-7)",
    ["Loaded by $1047. Generates the 40x40 map (river, hills, forest, end zones) from the game seed, derives",
     "the 5-character map id the two players compare, runs the unit setup screen and saves/loads game films.",
     "Shares its address range with the comcen screens overlay, which replaces it once the battle starts."], []),
 ("game/comcenScreens6F00", "B",  0x6F00, 0x8800, "$6F00-$87FF - comcen missile and drone screens (track 33 sectors 0-16 + track 32 sectors 8-15)",
    ["The variant the boot loader installs; $1086 reloads it. Implements the two command centre screens",
     "reached with the function keys: the missile screen ($7587) and the drone screen ($7FE6), including the",
     "radar rendering, plus the end-of-game result messages."], []),
 ("game/messages8800",   "B",     0x8800, 0x8A00, "$8800-$89FF - modem status message strings (track 29 sectors 9-10)",
    ["Only used while a modem game is being set up; the solo trainer AI reuses the whole block as work RAM."], []),
 ("game/graphicsData9300", "B",   0x9300, 0xA000, "$9300-$9FFF - screen address tables, fonts, terrain and unit graphics (track 28 sectors 0-12)", [], []),
 ("game/textEngineC000", "B",     0xC000, 0xD000, "$C000-$CFFF - text/glyph engine, status message system and unit pictures (track 30 sectors 0-15)",
    ["Loaded by $0F07 at start-up over the boot loader. Draws proportional text into the bitmap, manages the",
     "message queue and the status panel, and holds most of the game's vocabulary (unit names, group names,",
     "game types, menu text)."], []),
 ("game/trainerAiE000",  "B",     0xE000, 0xF000, "$E000-$EFFF - solo trainer: the computer opponent (track 35 sectors 0-15)",
    ["Loaded by $0F1A when the solo trainer was chosen ($0BA5 bit 7 set). It presents the same jump table as",
     "the modem driver but fakes the link: instead of sending and receiving over a wire it plans the moves of",
     "side 1 and injects them straight into the received-packet buffer. $EE00-$EFFF is swapped for the drone",
     "AI (track 29 sectors 4-5) while its drone is airborne."], [0xE000,0xE003,0xE006,0xE009,0xE00C,0xE012,0xE015]),
 ("game/modemDriverE000", "A",    0xE000, 0xF000, "$E000-$EFFF - modem driver: bit-banged RS-232, Hayes control and the packet protocol (track 18 sector 7 + track 34 sectors 1-15)",
    ["Loaded by $0F31 for a modem game ($0BA5 bit 7 clear). Three layers: a software UART on the user port",
     "driven from the CIA2 NMI (300/1200 baud), a Hayes AT command layer with a small terminal mode, and the",
     "packet layer that exchanges game commands and chat text with the opponent. Same jump table at $E000 as",
     "the solo trainer, so the rest of the game does not care which one is loaded."], [0xE000,0xE003,0xE006,0xE009,0xE00C,0xE012,0xE015]),
 ("game/trainerSetupEC00", "A2",  0xEC00, 0xF000, "$EC00-$EFFF - solo trainer setup: picks the opening play and seeds the AI (track 29 sectors 0-3)",
    ["Loaded by $7B24 from the map generator overlay, over the tail of the trainer module."], []),
 ("game/droneAiEE00",    "B3",    0xEE00, 0xF000, "$EE00-$EFFF - solo trainer: drone flight AI (track 29 sectors 4-5)",
    ["Swapped in by the trainer module ($E0D6) while its drone is in the air and swapped back out afterwards."], []),
 ("game/unitStartTemplatesF000", "B", 0xF000, 0xF700, "$F000-$F6FF - unit start-position templates (track 18 sectors 8-14)",
    ["What is STORED here is four 400-byte templates at $F000, $F190, $F320 and $F4B0; each is four parallel",
     "100-entry arrays - starting column, starting row, flags (facing in bits 0-1, side in bit 5, grouped in",
     "bit 6, group id in bits 2-4) and unit type. One template is chosen per game type and copied into the",
     "live unit arrays. The standard roster is 28 GRUNT, 12 RIDER, 6 BOOMER, 3 SPY and 1 COMCEN per side.",
     "What the same addresses hold at RUN TIME is different: $F000-$F63F becomes the 40x40 battlefield map",
     "(one byte per cell) and the unit record arrays start at $F640 in 100-byte slices, so the templates are",
     "overwritten as soon as a game starts.",
     "The 192 bytes at $F640-$F6FF are not used at all: they are a fragment of an older build of the game's",
     "own options-list printer, left in the sector by the development system."], []),
 ("game/highMemoryFBB8", "B",     0xFBB8, 0xFFD2, "$FBB8-$FFD1 - resident high memory: line walker, random numbers, arithmetic helpers",
    ["Track 29 sectors 11-15, loaded at $FA00 and then moved up by $1B8 so that it ends just below the CPU",
     "vectors. $FFDE-$FFE7 holds the 10-byte disk serial the two sides compare."], []),
]

# ---------------------------------------------------------------------------------------------
driveZpNames, driveZpComments = {}, {}
unitDataTypes = {}          # unit -> {addr: (kind, length, ...)} - overrides the global dataTypes

extraEntries = set()
unitLabels = {}          # unit -> {addr: name}

from buildAliases import UNIT_ALIASES



def canonicalUnit(u):
    return UNIT_ALIASES.get(u, u)


def loadAnnotations():
    labels, notes, dataTypes, zpNames, zpComments = {}, {}, {}, {}, {}
    for f in sorted(glob.glob(os.path.join(ANN, "*.json"))):
        j = json.load(open(f))
        for k, v in j.get("labels", {}).items():
            labels[int(k, 16)] = v
        for unit, m in j.get("unitLabels", {}).items():
            for k, v in m.items():
                unitLabels.setdefault(canonicalUnit(unit), {})[int(k, 16)] = v
        for k, v in j.get("notes", {}).items():
            a = int(k, 16)
            v = dict(v)
            for kk in ("block", "routine"):
                if isinstance(v.get(kk), str):
                    v[kk] = [v[kk]]
            unit = canonicalUnit(v.get("unit")) if v.get("unit") else None
            if unit:
                v["unit"] = unit
            lst = notes.setdefault(a, [])
            for cur in lst:
                if cur.get("unit") == unit:
                    cur.update(v)
                    break
            else:
                lst.append(v)
        fileUnit = canonicalUnit(j["unit"]) if j.get("unit") else None
        for k, v in j.get("dataTypes", {}).items():
            spec = tuple([v[0], int(v[1])] + [int(x) for x in v[2:]])
            if fileUnit:                      # a per-chunk annotation only speaks for its own file
                unitDataTypes.setdefault(fileUnit, {})[int(k, 16)] = spec
            else:
                dataTypes[int(k, 16)] = spec
        for unit, m in j.get("unitDataTypes", {}).items():
            for k, v in m.items():
                unitDataTypes.setdefault(canonicalUnit(unit), {})[int(k, 16)] = tuple([v[0], int(v[1])] + [int(x) for x in v[2:]])
        for k in j.get("entries", []):
            extraEntries.add(int(k, 16))
        for k, v in j.get("zp", {}).items():
            a = int(k, 16)
            if isinstance(v, list):
                zpNames[a] = v[0]
                if len(v) > 1: zpComments[a] = v[1]
            else:
                zpNames[a] = v
        for k, v in j.get("zpComments", {}).items():
            zpComments[int(k, 16)] = v
        for k, v in j.get("driveZp", {}).items():
            a = int(k, 16)
            if isinstance(v, list):
                driveZpNames[a] = v[0]
                if len(v) > 1: driveZpComments[a] = v[1]
            else:
                driveZpNames[a] = v
    return labels, notes, dataTypes, zpNames, zpComments


def main():
    os.makedirs(OUT, exist_ok=True)
    for f in glob.glob(os.path.join(OUT, "*", "*.s")) + glob.glob(os.path.join(OUT, "build", "*")):
        os.remove(f)
    userLabels, notes, dataTypes, zpNames, zpComments = loadAnnotations()
    hw = c64Symbols()

    # --- trace every runtime image from $0800 plus unit entry points (pass 1)
    OVERLAY_RANGES = [(0x6F00,0x8800),(0xE000,0xF000),(0xC000,0xD000),(0x0200,0x0400),(0x0400,0x0800)]
    def overlayOf(a):
        for lo, hi in OVERLAY_RANGES:
            if lo <= a < hi:
                return (lo, hi)
        return None
    imageCode = {}
    allTargets = {}           # addr -> kind ("sub"|"loc"|"data")
    apiEntries = set()        # targets inside an overlay range that are reached from outside that range
    unitEntries = set(extraEntries)
    for u in UNITS:
        if u[1] in IMAGES:
            unitEntries |= set(u[6])
    for key, img in IMAGES.items():
        code, calls, jumps, ind, ill = trace(img, {0x0800} | unitEntries, validFor(RUNTIME_VALID))
        for t, sources in list(calls.items()) + list(jumps.items()):
            ov = overlayOf(t)
            if ov and any(not (ov[0] <= s < ov[1]) for s in sources):
                apiEntries.add(t)
    # pass 2: natural entries first, then each API entry transactionally (kept only if its path is clean)
    def traceClean(img, baseEntries, extraEntries):
        """Trace natural non-overlay entries freely; every other entry (overlay API entries, discovered
        vectors/tables, annotation entries) is traced in a transaction confined to (non-overlay + its own
        overlay) and kept only if no illegal opcode is reached."""
        def validOutsideOverlays(a):
            return validFor(RUNTIME_VALID)(a) and overlayOf(a) is None
        def validWithin(ov):
            if ov is None:
                return validOutsideOverlays
            return lambda a: validFor(RUNTIME_VALID)(a) and (overlayOf(a) is None or overlayOf(a) == ov)
        code = set()
        allCalls, allJumps, allInd, allIll = {}, {}, [], []
        rejected = []
        done = set()
        candidates = set()
        def absorb(c2, calls2, jumps2, ind2):
            code.update(c2)
            for t, s in calls2.items():
                allCalls.setdefault(t, []).extend(s)
                if overlayOf(t): candidates.add(t)
            for t, s in jumps2.items():
                allJumps.setdefault(t, []).extend(s)
                if overlayOf(t): candidates.add(t)
            allInd.extend(ind2)
        base = {e for e in baseEntries if overlayOf(e) is None}
        c2, calls2, jumps2, ind2, ill2 = trace(img, base, validOutsideOverlays)
        allIll.extend(ill2)
        absorb(c2, calls2, jumps2, ind2)
        candidates |= {e for e in baseEntries if overlayOf(e)} | set(extraEntries)
        while candidates - done:
            e = min(candidates - done)
            done.add(e)
            if e in code:
                continue
            c2, calls2, jumps2, ind2, ill2 = trace(img, {e}, validWithin(overlayOf(e)), known=code)
            if ill2:
                rejected.append(e)
                continue
            absorb(c2, calls2, jumps2, ind2)
        return code, allCalls, allJumps, allInd, allIll, rejected

    surveyCodeEntries = set()
    for sf in glob.glob(os.path.join(ROOT, "survey", "*.json")):
        try:
            sj = json.load(open(sf))
        except Exception:
            continue
        for m in sj.get("misclassified", []):
            if m.get("actual") == "code":
                try:
                    surveyCodeEntries.add(int(str(m["addr"]).replace("$", ""), 16))
                except Exception:
                    pass
    discoveredTables = {}        # base -> number of words (rendered as .addr)
    splitTables = {}             # low-byte table base -> (high-byte table base, entry count)
    def traceWithDiscovery(img, key, baseEntries, extraEntries):
        tried = set()
        extra = set(extraEntries)
        def plausible(w, code, starts):
            if not validFor(RUNTIME_VALID)(w) or w < 0x0200:
                return False
            if w in code and w not in starts:
                return False                      # inside an existing instruction
            i = decode(img, w)
            return i.isLegal() and i.mnemonic != "brk"
        for _ in range(8):
            code, calls, jumps, ind, ill, rejected = traceClean(img, baseEntries, extra)
            starts = instructionStarts(img, code)
            pl = lambda w: plausible(w, code, starts)
            vecs = {0xFFFA, 0xFFFE} | {decode(img, a).operand for a in ind}
            cands = set(vectorEntries(img, starts, vecs))
            for kind, baseLo, baseHi in patchedCallTables(img, starts):
                if kind == "word":
                    ws = wordsFrom(img, baseLo, pl)
                    if ws:
                        discoveredTables[baseLo] = max(discoveredTables.get(baseLo, 0), len(ws))
                        cands |= set(ws)
                else:
                    ws = splitTableEntries(img, baseLo, baseHi, pl)
                    cands |= set(ws)
                    splitTables[baseLo] = (baseHi, max(len(ws), splitTables.get(baseLo, (0, 0))[1]))
            for e in surveyCodeEntries:
                cands.add(e)
            new = {c for c in cands if pl(c)} - tried - code
            if not new:
                break
            tried |= new
            extra |= new
        return code, calls, jumps, ind, ill, rejected

    for key, img in IMAGES.items():
        code, calls, jumps, ind, ill, rejected = traceWithDiscovery(img, key, {0x0800} | unitEntries, apiEntries)
        imageCode[key] = code
        for t in calls: allTargets[t] = "sub"
        for t in jumps: allTargets.setdefault(t, "loc")
        for a in ill:
            print(f"[{key}] warning: traced into illegal opcode ${img[a]:02X} at ${a:04X}")
        if rejected:
            shown = " ".join(f"${a:04X}" for a in rejected[:12]) + (" ..." if len(rejected) > 12 else "")
            print(f"[{key}] rejected {len(rejected)} candidate entries (not clean code in this image): {shown}")
        for a in ind:
            i = decode(img, a)
            print(f"[{key}] note: indirect jmp (${i.operand:04X}) at ${a:04X}")

    # discovered dispatch tables render as .addr unless an annotation says otherwise
    for b, n in discoveredTables.items():
        if b not in dataTypes and n >= 2:
            dataTypes[b] = ("addr", 2 * n)

    # boot-time images
    bootEa = bytearray(65536)
    bootEa[0x02A8:0x02A8+len(eaPrg)-2] = eaPrg[2:]
    bootEa[0xC000:0xC400] = bootImage[0xC000:0xC400]
    code, calls, jumps, ind, ill = trace(bootEa, {0x02B8}, validFor([(0x02A8,0x030C)]))
    imageCode["BOOTEA"] = code
    for t in calls: allTargets[t] = "sub"
    for t in jumps: allTargets.setdefault(t, "loc")
    code, calls, jumps, ind, ill = trace(bootImage, {0xC145, 0xC004, 0xC034, 0xC229, 0xC3C4, 0xC03B, 0xC08A}, validFor([(0xC000,0xC400)]))
    imageCode["BOOT"] = code
    for t in calls: allTargets[t] = "sub"
    for t in jumps: allTargets.setdefault(t, "loc")
    IMAGES["BOOTEA"] = bootEa
    IMAGES["BOOT"]   = bootImage
    loadImg = bytearray(65536)
    loadImg[0x9800:0x9800+len(loadPrg)] = loadPrg
    IMAGES["LOADPRG"] = loadImg
    imageCode["LOADPRG"] = set()

    # branch targets and data references become labels too
    externalRefs = set()      # targets referenced from outside their own overlay range
    def noteRef(src, t):
        ov = overlayOf(t)
        if ov is None or not (ov[0] <= src < ov[1]):
            externalRefs.add(t)
    def collectRefs(img, codeSet, validRanges):
        addrs = sorted(codeSet)
        i = 0
        while i < len(addrs):
            a = addrs[i]
            insn = decode(img, a)
            if not insn.isLegal():
                i += 1; continue
            if insn.mode == REL:
                allTargets.setdefault(insn.target, "loc")
                noteRef(a, insn.target)
            elif insn.mode in (ZP, ZPX, ZPY, IZX, IZY):
                pass
            elif insn.mode in (ABS, ABX, ABY, IND):
                t = insn.target
                noteRef(a, t)
                if insn.mnemonic not in ("jmp", "jsr") and t >= 0x100 and t not in hw and (0xD000 > t or t >= 0xE000):
                    allTargets.setdefault(t, "data")
            # skip operand bytes
            j = i + 1
            while j < len(addrs) and addrs[j] < a + insn.length:
                j += 1
            i = j
    for key, img in IMAGES.items():
        collectRefs(img, imageCode[key], None)

    # --- global labels
    labels = {}
    for t, kind in allTargets.items():
        if t in hw: continue
        if kind == "sub":   labels[t] = f"sub_{t:04X}"
        elif kind == "loc": labels[t] = f"L_{t:04X}"
        else:               labels[t] = f"D_{t:04X}"
    labels.update(userLabels)
    # zero page names: default zp_XX for every referenced zp address
    zpUsed = set()
    for key, img in IMAGES.items():
        for a in imageCode[key]:
            insn = decode(img, a)
            if insn.isLegal() and insn.mode in (ZP, ZPX, ZPY, IZX, IZY):
                zpUsed.add(insn.operand)
            if insn.isLegal() and insn.mode in (ABS, ABX, ABY) and insn.operand < 0x100:
                zpUsed.add(insn.operand)
    allZp = {a: zpNames.get(a, f"zp_{a:02X}") for a in sorted(zpUsed | set(zpNames))}
    seenZp = {}
    for a in sorted(allZp):
        n = allZp[a]
        if n in seenZp:
            print(f"warning: zero-page name '{n}' used for ${seenZp[n]:02X} and ${a:02X}; renaming the latter")
            allZp[a] = f"{n}_{a:02X}"
        seenZp[allZp[a]] = a
    writeIncludes(OUT, allZp, zpComments)

    # --- render units
    index = []
    verify = ["#!/bin/bash", "# verify.sh - reassemble every source file with ca65 and compare against the original bytes.",
              "set -e", "cd \"$(dirname \"$0\")\"", "mkdir -p build", "fail=0"]
    # which units cover an address, and whether it is an instruction start / labelled there
    def insnStartsOf(key):
        starts = set()
        img = IMAGES[key]
        nextStart = 0
        for a in sorted(imageCode[key]):
            if a < nextStart: continue
            insn = decode(img, a)
            if insn.isLegal():
                starts.add(a)
                nextStart = a + insn.length
        return starts
    startsCache = {key: insnStartsOf(key) for key in IMAGES}
    unitInfo = {name: (key, start, end) for name, key, start, end, *_ in UNITS}
    variantNames = {}                        # addr -> {unit: name}
    for unit, m in unitLabels.items():
        for a, n in m.items():
            variantNames.setdefault(a, {})[unit] = n
    def isRealThere(addr, unit):
        key, s, e = unitInfo.get(unit, (None, 0, 0))
        if key is None or not (s <= addr < e):
            return False
        if addr in startsCache.get(key, ()):
            return True
        return addr not in imageCode[key]        # data byte (not inside an instruction)
    def externalName(addr):
        cands = [(u, n) for u, n in variantNames.get(addr, {}).items() if isRealThere(addr, u)]
        if not cands:
            cands = list(variantNames.get(addr, {}).items())
        if not cands:
            return None, None
        cands.sort()
        chosen = cands[0][1]
        alt = None
        if len(cands) > 1:
            alt = "also " + ", ".join(f"{u.split('/')[-1]}:{n}" for u, n in cands[1:])
        return chosen, alt
    def uniquify(merged, unit):
        seen = {}
        for a in sorted(merged):
            n = merged[a]
            if n in seen and seen[n] != a:
                newName = f"{n}_{a:04X}"
                print(f"[{unit}] warning: label name '{n}' used for ${seen[n]:04X} and ${a:04X}; renaming the latter to {newName}")
                merged[a] = newName
                n = newName
            seen[n] = a
        for n in set(merged.values()):
            if n in hw.values() or n.lower() in ("a", "x", "y"):
                for a in [k for k, v in merged.items() if v == n]:
                    merged[a] = f"{n}_{a:04X}"
                    print(f"[{unit}] warning: label name '{n}' clashes with a reserved/hardware name; renamed")
        return merged

    def labelsForUnit(unit):
        merged = dict(labels)
        alts = {}
        key, uStart, uEnd = unitInfo.get(unit, (None, 0, 0))
        for a, m in variantNames.items():
            if unit in m:
                merged[a] = m[unit]
            elif uStart <= a < uEnd:
                continue                      # another variant's name for an address inside this file
            else:
                n, alt = externalName(a)
                if n:
                    merged[a] = n
                if alt:
                    alts[a] = alt
        return uniquify(merged, unit), alts
    def notesForUnit(unit):
        out = {}
        for a, lst in notes.items():
            merged = {}
            for v in lst:                       # global first, unit-specific overrides
                if v.get("unit") is None:
                    merged.update(v)
            for v in lst:
                if v.get("unit") == unit:
                    merged.update(v)
            if merged:
                out[a] = merged
        return out

    runtimeKeys = [k for k in IMAGES if k not in ("BOOT", "BOOTEA", "LOADPRG")]
    unionCode = set()
    for k in runtimeKeys:
        unionCode |= imageCode[k]
    for k in runtimeKeys:
        # non-overlay ranges hold identical bytes in every runtime image: share the union of traced code
        extra = {a for a in unionCode if overlayOf(a) is None}
        imageCode[k] = imageCode[k] | extra
    startsCache = {key: insnStartsOf(key) for key in IMAGES}

    for name, key, start, end, title, prov, extra in UNITS:
        img = IMAGES[key]
        codeSet = imageCode[key]
        isBoot = name.startswith("boot/")
        hwHere = hw if isBoot else {a: n for a, n in hw.items() if a not in KERNAL}
        uLabels, uAlts = labelsForUnit(name)
        # a hardware/KERNAL name always wins over a game label at the same address: in the boot units
        # $FFxx is the KERNAL jump table, while at run time the same addresses are ordinary RAM
        for a in list(uLabels):
            if a in hwHere:
                del uLabels[a]
        dtHere = dict(dataTypes)
        dtHere.update(unitDataTypes.get(name, {}))
        r = RendererT(img, codeSet, uLabels, notesForUnit(name), hwHere, dtHere, allZp, useKernal=isBoot)
        r.altComments = uAlts
        header = list(prov)
        externs = []
        lines, defined = r.renderRegion(start, end, title, header, externs)
        path = os.path.join(OUT, name + ".s")
        os.makedirs(os.path.dirname(path), exist_ok=True)
        open(path, "w").write("\n".join(lines) + "\n")
        binPath = os.path.join(OUT, "build", name.replace("/", "_") + ".orig.bin")
        os.makedirs(os.path.dirname(binPath), exist_ok=True)
        open(binPath, "wb").write(img[start:end])
        base = name.replace("/", "_")
        verify.append(f"ca65 -t none -I . -o build/{base}.o {name}.s && ld65 -C build/link.cfg -o build/{base}.bin build/{base}.o && "
                      f"(cmp -s build/{base}.bin build/{base}.orig.bin && echo 'OK   {name}' || {{ echo 'FAIL {name}'; fail=1; }})")
        codeBytes = sum(1 for a in range(start, end) if a in codeSet)
        index.append((name, start, end, codeBytes, title))
    # --- 1541 drive-side code (plain, unencrypted sectors on track 1)
    driveZpLines = ["; drive_zp.inc - 1541 DOS zero-page locations used by the EA drive code", ""]
    for a in range(0x100):
        n = driveZpNames.get(a, f"zp_{a:02X}")
        c = driveZpComments.get(a, "")
        driveZpLines.append(f"{n:<24} := ${a:02X}" + (f"    ; {c}" if c else ""))
    open(os.path.join(OUT, "drive_zp.inc"), "w").write("\n".join(driveZpLines) + "\n")
    driveLines = ["; drive1541.inc - 1541 VIA registers and DOS ROM entry points used by the loader", ""]
    for a in sorted(DRIVE):
        driveLines.append(f"{DRIVE[a]:<24} := ${a:04X}")
    open(os.path.join(OUT, "drive1541.inc"), "w").write("\n".join(driveLines) + "\n")
    DRIVE_UNITS = [
      ("drive/driveBootstrap", rawSec(1,17), 0x0600, [0x0600, 0x0624],
       "1541 bootstrap sector (track 1 sector 17) - executed in the drive by DOS command \"B-E 2 0 1 17\"",
       ["DOS loads this sector into the buffer of channel 2 and jumps to its first byte.  The code finds its own",
        "page with a JSR/TSX trick, copies itself to buffer 3 ($0600) and continues there at $0624, so it is shown",
        "here at its final address.  It queues three job-code reads (track 1 sectors 18,19,20 -> $0300-$05FF),",
        "re-initialises the drive and jumps to $0300 (the fast loader).  The tail of the sector holds a 1988",
        "Electronic Arts message that is never displayed."]),
      ("drive/driveFastLoader", rawSec(1,18) + rawSec(1,19) + rawSec(1,20), 0x0300, [0x0300],
       "1541 fast loader (track 1 sectors 18-20, runs at $0300-$05FF in the drive)",
       ["Commands from the C64 (received through the 2-bit protocol on the serial bus):",
        "  $C0 = reset drive, $60 = write sector (receive 256 bytes, encrypt, write), anything else = read sector.",
        "Reads are done with the DOS ROM helpers (header search, GCR decode, checksum) into buffer 3 ($0600),",
        "decrypted with the rolling XOR key at $0531 (not applied to track 18) and sent back 2 bits at a time."]),
    ]
    t1s0 = rawSec(1, 0)
    RAW_UNITS = [
      ("boot/c128BootSector", t1s0, 0x0B00, [0x0B27],
       "C128 auto-boot sector (track 1 sector 0, loaded by the C128 KERNAL to $0B00)",
       ["Header: \"CBM\", load address 0, bank 0, 0 extra blocks, boot message \"AN ELECTRONIC ARTS PRODUCTION\".",
        "The code copies $0B89-$0BFF to $8000 (a CBM80 cartridge image), puts a 1571 into 1541 mode with",
        "\"U0>M0\" / \"U0>H0\" on the command channel, selects MMU bank 15 and calls GO64 ($FF4D).  In C64 mode",
        "the cartridge signature at $8004 autostarts the copied code (see boot/c64CartridgeStub8000.s)."], "c64"),
      ("boot/c64CartridgeStub8000", t1s0[0x89:], 0x8000, [0x8009],
       "Pseudo-cartridge image copied to $8000 by the C128 boot sector (bytes $89-$FF of track 1 sector 0)",
       ["$8000/$8002 = cold/warm start vectors ($8009), $8004 = \"CBM80\".  The C64 KERNAL reset code jumps to",
        "$8009, which initialises I/O, clears low RAM, sets MEMTOP and LOADs the normal boot file."], "c64"),
    ]
    for name, data, base, entries, title, prov, symSet in RAW_UNITS:
        img = bytearray(65536)
        img[base:base+len(data)] = data
        end = base + len(data)
        code, calls, jumps, ind, ill = trace(img, set(entries), validFor([(base, end)]))
        uLabels = {}
        for t in calls:
            if t not in hw: uLabels[t] = f"sub_{t:04X}"
        for t in jumps:
            if t not in hw: uLabels.setdefault(t, f"L_{t:04X}")
        addrs = sorted(code)
        i = 0
        while i < len(addrs):
            a = addrs[i]
            insn = decode(img, a)
            if insn.isLegal():
                if insn.mode == REL:
                    uLabels.setdefault(insn.target, f"L_{insn.target:04X}")
                elif insn.mode in (ABS, ABX, ABY) and insn.mnemonic not in ("jmp","jsr") and base <= insn.target < end:
                    uLabels.setdefault(insn.target, f"D_{insn.target:04X}")
                j = i + 1
                while j < len(addrs) and addrs[j] < a + insn.length:
                    j += 1
                i = j
            else:
                i += 1
        uLabels.update(unitLabels.get(name, {}))
        uNotes = notesForUnit(name)
        uTypes = {a: v for a, v in dataTypes.items() if base <= a < end}
        uTypes.update({a: v for a, v in unitDataTypes.get(name, {}).items() if base <= a < end})
        r = RendererT(img, code, uLabels, uNotes, hw, uTypes, allZp, useKernal=True)
        lines, defined = r.renderRegion(base, end, title, prov, [])
        path = os.path.join(OUT, name + ".s")
        open(path, "w").write("\n".join(lines) + "\n")
        open(os.path.join(OUT, "build", name.replace("/", "_") + ".orig.bin"), "wb").write(data)
        b = name.replace("/", "_")
        verify.append(f"ca65 -t none -I . -o build/{b}.o {name}.s && ld65 -C build/link.cfg -o build/{b}.bin build/{b}.o && "
                      f"(cmp -s build/{b}.bin build/{b}.orig.bin && echo 'OK   {name}' || {{ echo 'FAIL {name}'; fail=1; }})")
        index.append((name, base, end, len(code), title))

    for name, data, base, entries, title, prov in DRIVE_UNITS:
        img = bytearray(65536)
        img[base:base+len(data)] = data
        end = base + len(data)
        code, calls, jumps, ind, ill = trace(img, set(entries), validFor([(base, end)]))
        dLabels = {}
        for t in calls:
            if t not in DRIVE: dLabels[t] = f"sub_{t:04X}"
        for t in jumps:
            if t not in DRIVE: dLabels.setdefault(t, f"L_{t:04X}")
        addrs = sorted(code)
        i = 0
        while i < len(addrs):
            a = addrs[i]
            insn = decode(img, a)
            if insn.isLegal():
                if insn.mode == REL:
                    dLabels.setdefault(insn.target, f"L_{insn.target:04X}")
                elif insn.mode in (ABS, ABX, ABY) and insn.mnemonic not in ("jmp","jsr") and base <= insn.target < end:
                    dLabels.setdefault(insn.target, f"D_{insn.target:04X}")
                j = i + 1
                while j < len(addrs) and addrs[j] < a + insn.length:
                    j += 1
                i = j
            else:
                i += 1
        dLabels.update(unitLabels.get(name, {}))
        dNotes = notesForUnit(name)
        dTypes = {a: v for a, v in dataTypes.items() if base <= a < end}
        dTypes.update({a: v for a, v in unitDataTypes.get(name, {}).items() if base <= a < end})
        zpD = {a: f"zp_{a:02X}" for a in range(0x100)}
        zpD.update(driveZpNames)
        r = RendererT(img, code, dLabels, dNotes, DRIVE, dTypes, zpD, useKernal=False)
        lines, defined = r.renderRegion(base, end, title, prov, [])
        lines = [l.replace('.include "c64.inc"', '.include "drive1541.inc"').replace('.include "zeropage.inc"', '.include "drive_zp.inc"') for l in lines]
        path = os.path.join(OUT, name + ".s")
        os.makedirs(os.path.dirname(path), exist_ok=True)
        open(path, "w").write("\n".join(lines) + "\n")
        open(os.path.join(OUT, "build", name.replace("/", "_") + ".orig.bin"), "wb").write(data)
        b = name.replace("/", "_")
        verify.append(f"ca65 -t none -I . -o build/{b}.o {name}.s && ld65 -C build/link.cfg -o build/{b}.bin build/{b}.o && "
                      f"(cmp -s build/{b}.bin build/{b}.orig.bin && echo 'OK   {name}' || {{ echo 'FAIL {name}'; fail=1; }})")
        index.append((name, base, end, len(code), title))

    # --- cross reference
    unitOf = []
    for name, key, start, end, *_ in UNITS:
        unitOf.append((start, end, key, name))
    def unitName(a, key):
        for s, e, k, n in unitOf:
            if s <= a < e and k == key:
                return n
        for s, e, k, n in unitOf:
            if s <= a < e:
                return n
        return "?"
    refs = {}      # target -> set of (kind, srcAddr, unit)
    zpRefs = {}
    for key, img in IMAGES.items():
        addrs = sorted(imageCode[key])
        i = 0
        while i < len(addrs):
            a = addrs[i]
            insn = decode(img, a)
            if not insn.isLegal():
                i += 1; continue
            if insn.mode in (ZP, ZPX, ZPY, IZX, IZY):
                kind = "w" if insn.mnemonic in ("sta","stx","sty","inc","dec","asl","lsr","rol","ror") else "r"
                zpRefs.setdefault(insn.operand, set()).add((kind, a, unitName(a, key)))
            elif insn.mode in (ABS, ABX, ABY, IND, REL):
                t = insn.target
                if insn.mnemonic == "jsr": kind = "call"
                elif insn.mnemonic == "jmp": kind = "jmp"
                elif insn.mode == REL: kind = "br"
                elif insn.mnemonic in ("sta","stx","sty","inc","dec","asl","lsr","rol","ror"): kind = "w"
                else: kind = "r"
                if t < 0x100:
                    zpRefs.setdefault(t, set()).add((kind, a, unitName(a, key)))
                else:
                    refs.setdefault(t, set()).add((kind, a, unitName(a, key)))
            j = i + 1
            while j < len(addrs) and addrs[j] < a + insn.length:
                j += 1
            i = j
    def unitsContaining(a):
        return [n for s, e, k, n in unitOf if s <= a < e]

    def formatRefs(entries):
        byKind = {}
        for kind, addr in sorted(entries, key=lambda x: x[1]):
            byKind.setdefault(kind, []).append(f"{addr:04X}")
        return " ".join(f"{k}:" + ",".join(v[:40]) + ("..." if len(v) > 40 else "") for k, v in byKind.items())

    with open(os.path.join(OUT, "XREF.txt"), "w") as f:
        f.write("; Cross reference for the Modem Wars disassembly.\n")
        f.write("; Each entry is:   name  $address  [source files that contain this address]\n")
        f.write(";                    from <file>: kind:addresses ...\n")
        f.write("; kinds: call = JSR, jmp = JMP, br = branch, r = read, w = write.\n")
        f.write("; IMPORTANT: $6F00-$87FF and $E000-$EFFF each exist in two different files; a reference is\n")
        f.write("; only meaningful for the file named on its 'from' line.\n\n")
        for t in sorted(refs):
            n = (externalName(t)[0] if t in variantNames else None) or labels.get(t) or hw.get(t) or f"${t:04X}"
            where = ",".join(unitsContaining(t)) or "-"
            f.write(f"{n:<28} ${t:04X} [{where}]\n")
            bySource = {}
            for kind, addr, srcUnit in refs[t]:
                bySource.setdefault(srcUnit, []).append((kind, addr))
            for srcUnit in sorted(bySource):
                f.write(f"    from {srcUnit}: {formatRefs(bySource[srcUnit])}\n")
        f.write("\n; Zero page usage\n\n")
        for z in sorted(zpRefs):
            n = allZp.get(z, f"${z:02X}")
            f.write(f"{n:<28} ${z:02X}\n")
            bySource = {}
            for kind, addr, srcUnit in zpRefs[z]:
                bySource.setdefault(srcUnit, []).append((kind, addr))
            for srcUnit in sorted(bySource):
                f.write(f"    from {srcUnit}: {formatRefs(bySource[srcUnit])}\n")

    open(os.path.join(OUT, "build", "link.cfg"), "w").write(
        "MEMORY { RAM: start = $0000, size = $10000, file = %O; }\nSEGMENTS { CODE: load = RAM, type = rw; }\n")
    verify.append("exit $fail")
    open(os.path.join(OUT, "verify.sh"), "w").write("\n".join(verify) + "\n")
    os.chmod(os.path.join(OUT, "verify.sh"), 0o755)
    with open(os.path.join(OUT, "INDEX.txt"), "w") as f:
        for name, s, e, cb, t in index:
            f.write(f"{name:<22} ${s:04X}-${e-1:04X}  code {cb:5d}/{e-s:5d}  {t}\n")
    print(open(os.path.join(OUT, "INDEX.txt")).read())
    json.dump({f"{a:04X}": n for a, n in sorted(labels.items())}, open(os.path.join(OUT, "build", "labels.json"), "w"), indent=0)
    # authoritative per-unit map of the byte ranges that are NOT code (used by tools/findStrings.py)
    dataMap = {}
    for name, key, start, end, *_ in UNITS:
        codeSet = imageCode[key]
        ranges = []
        a = start
        while a < end:
            if a in codeSet:
                a += 1
                continue
            b2 = a
            while b2 < end and b2 not in codeSet:
                b2 += 1
            ranges.append([f"{a:04X}", f"{b2 - 1:04X}"])
            a = b2
        dataMap[name] = ranges
    json.dump(dataMap, open(os.path.join(OUT, "build", "dataMap.json"), "w"), indent=0)

if __name__ == "__main__":
    main()
