#!/usr/bin/env python3
# chunks.py - split the rendered listings into agent-sized chunks at routine boundaries.
# Writes disassembly/build/chunks.json : [{id, file, unit, startAddr, endAddr, startLine, endLine, codeLines}]
import os, re, json, sys
ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
OUT = os.path.join(ROOT, "disassembly")
TARGET = int(sys.argv[1]) if len(sys.argv) > 1 else 700      # target code lines per chunk

ADDR_RE = re.compile(r";\s([0-9A-F]{4})(\s|$)")

def chunkFile(relPath):
    lines = open(os.path.join(OUT, relPath)).read().split("\n")
    # candidate cut points: a line that is a label definition for a sub_/named routine, whose previous
    # non-empty line is rts/rti/jmp or a data row or a comment banner
    insnCount = 0
    cuts = []
    addrAt = {}
    lastAddr = None
    for i, l in enumerate(lines):
        m = None if l.startswith(";") else ADDR_RE.search(l)   # ignore comment-only lines
        if m:
            lastAddr = int(m.group(1), 16)
            addrAt[i] = lastAddr
        if re.match(r"^[A-Za-z_][A-Za-z0-9_]*:\s*$", l) and i > 0:
            j = i - 1
            while j > 0 and (lines[j].strip() == "" or lines[j].startswith(";")):
                j -= 1
            prev = lines[j].strip()
            if re.match(r"^(rts|rti|jmp)\b", prev) or prev.startswith(".byte") or prev.startswith(".word") or prev.startswith(".addr"):
                cuts.append(i)
    orgLine = next(i for i, l in enumerate(lines) if l.strip().startswith(".org"))
    codeLineIdx = [i for i in range(len(lines)) if re.match(r"^\s+[a-z]{3}\b", lines[i])]
    chunks = []
    start = orgLine + 1
    count = 0
    lastCut = start
    for i in range(start, len(lines)):
        if re.match(r"^\s+[a-z]{3}\b", lines[i]):
            count += 1
        if i in cuts and count >= TARGET:
            chunks.append((lastCut, i - 1))
            lastCut = i
            count = 0
    chunks.append((lastCut, len(lines) - 1))
    out = []
    for s, e in chunks:
        addrs = [addrAt[i] for i in range(s, e + 1) if i in addrAt]
        if not addrs:
            continue
        n = sum(1 for i in range(s, e + 1) if re.match(r"^\s+[a-z]{3}\b", lines[i]))
        out.append({"file": relPath, "startAddr": f"{min(addrs):04X}", "endAddr": f"{max(addrs):04X}",
                    "startLine": s + 1, "endLine": e + 1, "codeLines": n})
    return out

def main():
    files = []
    for d in ("boot", "drive", "game"):
        for f in sorted(os.listdir(os.path.join(OUT, d))):
            if f.endswith(".s"):
                files.append(f"{d}/{f}")
    allChunks = []
    for f in files:
        for c in chunkFile(f):
            c["id"] = f"{f.replace('/', '_')[:-2]}_{c['startAddr']}"
            c["unit"] = f[:-2]
            allChunks.append(c)
    json.dump(allChunks, open(os.path.join(OUT, "build", "chunks.json"), "w"), indent=1)
    for c in allChunks:
        print(f"{c['id']:<34} lines {c['startLine']:>5}-{c['endLine']:>5}  ${c['startAddr']}-${c['endAddr']}  code lines {c['codeLines']}")
    print(len(allChunks), "chunks")

main()