#!/usr/bin/env python3
"""Find the data blocks in a project's prg.s that are really code, and list
their start addresses as tracer seeds.

A block that decodes as a long run of instructions was never reached by the
tracer - it sits behind a dispatch the trace could not follow - and it is
code, not data.  Seeding the tracer at the block's first byte makes it
reach the block and classify it correctly, which is what turns a "data"
block holding a routine back into that routine.

Usage:
    python3 _tools/code_data_seeds.py <project dir> [--min 4]
"""

import argparse
import re
import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parent))
import nesdis  # noqa: E402

BLOCK = re.compile(r"^; ---- data \$([0-9A-F]{4})-\$([0-9A-F]{4}) \((\d+) bytes\) ----$")
BYTE = re.compile(r"^    \.byte (.*?)\s*; CPU:\$([0-9A-F]{4}) PRG:[0-9A-F]{4}\s+((?:[0-9A-F]{2} ?)+)")


def main():
    ap = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    ap.add_argument("project")
    ap.add_argument("--min", type=int, default=4,
                    help="the shortest run of instructions that counts as code")
    args = ap.parse_args()

    d = Path(args.project)
    prg = (d / "prg.bin").read_bytes()
    base = 0x8000
    src = (d / "prg.s").read_text()
    m = re.search(r"PRG mapped at \$([0-9A-F]{4})", src)
    if m:
        base = int(m.group(1), 16)

    lines = src.split("\n")
    seeds = []
    i = 0
    while i < len(lines):
        b = BLOCK.match(lines[i])
        if not b:
            i += 1
            continue
        lo, hi = int(b.group(1), 16), int(b.group(2), 16)
        data = []
        j = i + 1
        while j < len(lines):
            r = BYTE.match(lines[j])
            if not r:
                break
            data += [int(t.strip().lstrip("$"), 16)
                     for t in r.group(1).split(",") if t.strip()]
            j += 1
        # decode
        off, count, first_bad = 0, 0, None
        while off < len(data):
            ent = nesdis.OPCODES.get(data[off])
            if ent is None:
                first_bad = off
                break
            n = nesdis.MODE_LEN[ent[1]]
            if off + n > len(data):
                first_bad = off
                break
            off += n
            count += 1
        if count >= args.min:
            seeds.append((lo, count, len(data)))
        i = j

    print(f"# {len(seeds)} data block(s) that decode as code")
    for lo, count, size in seeds:
        print(f"#  ${lo:04X}: {count} instructions in {size} bytes")
    print()
    print(",".join(f"0x{lo:X}" for lo, _c, _s in seeds))


if __name__ == "__main__":
    main()