#!/usr/bin/env python3
"""Report what each data block in a project really is, with the evidence.

The tracer already leaves the evidence in `prg.s`: every `; ---- data
$XXXX-$YYYY` block is followed by the lines

    ;   referenced by `lda $80D9,x` at $8081 (operand $80D9 falls inside this region)

one per instruction whose operand lands inside the block.  This tool collects
those, and adds what can be worked out from the bytes themselves:

  * whether the block decodes cleanly as instructions (so it is code the
    tracer could not reach, not data),
  * whether it looks like a table of 16-bit pointers into the PRG,
  * whether it looks like a table of 8-bit values in a narrow range,
  * whether it is all one repeated byte (padding).

The point is that a data block gets a name from *this* evidence, not from a
guess about what the game's data usually looks like.

Usage:
    python3 _tools/analyze_data_blocks.py <project dir> [--all]
"""

import argparse
import re
import sys
from collections import Counter
from pathlib import Path

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

BLOCK_RE = re.compile(
    r"^; ---- data \$([0-9A-F]{4})-\$([0-9A-F]{4}) \((\d+) bytes\) ----$")
REF_RE = re.compile(
    r"^;   referenced by `([^`]+)` at \$([0-9A-F]{4})")
BYTE_ROW_RE = re.compile(
    r"^    \.byte ((?:\$[0-9A-F]{2}(?:, )?)+)\s*; CPU:\$([0-9A-F]{4}) ")
LABEL_RE = re.compile(r"^([A-Za-z_][A-Za-z0-9_]*):\s*$")


def scan_consumers(path, blocks):
    """Find the instructions whose operand lands in a block.

    The tracer's own ";   referenced by ..." lines only cover the references
    it noticed while tracing; this scans the emitted source for every
    absolute operand and records which block it lands in.  The operand, the
    instruction and the address it sits at are all kept, because that is the
    evidence a table's name has to be based on.
    """
    insn_re = re.compile(
        r"^    ([A-Za-z@][A-Za-z0-9_]*)\s+(\S+)\s*; CPU:\$([0-9A-F]{4}) ")
    abs_re = re.compile(r"^(?:a:)?\$([0-9A-F]{4})$")
    for line in path.read_text().split("\n"):
        m = insn_re.match(line)
        if not m:
            continue
        mn, op, at = m.group(1), m.group(2), int(m.group(3), 16)
        a = abs_re.match(op)
        if not a:
            continue
        target = int(a.group(1), 16)
        for blk in blocks:
            lo, hi = blk[0], blk[1]
            if lo <= target <= hi:
                blk[4].append((f"{mn} {op}", at))
                break


def collect(path):
    """[(lo, hi, size, label, [(operand, at), ...], bytes)]"""
    lines = path.read_text().split("\n")
    blocks = []
    i = 0
    while i < len(lines):
        m = BLOCK_RE.match(lines[i])
        if not m:
            i += 1
            continue
        lo, hi, size = int(m.group(1), 16), int(m.group(2), 16), int(m.group(3))
        refs = []
        j = i + 1
        while j < len(lines) and lines[j].startswith(";   "):
            r = REF_RE.match(lines[j])
            if r:
                refs.append((r.group(1), int(r.group(2), 16)))
            j += 1
        label = None
        if j < len(lines):
            lm = LABEL_RE.match(lines[j])
            if lm:
                label = lm.group(1)
        # the bytes
        data = []
        k = j
        while k < len(lines):
            b = BYTE_ROW_RE.match(lines[k])
            if not b:
                # A label or a comment may sit between the block header and
                # its first `.byte` row; skip those, but stop at anything
                # else so the block does not run into the next one.
                if LABEL_RE.match(lines[k]) or lines[k].startswith(";"):
                    # Skip a label or a comment, but never step over the
                    # start of the next block.
                    if lines[k].startswith("; ---- data"):
                        break
                    k += 1
                    continue
                break
            for tok in b.group(1).split(","):
                data.append(int(tok.strip().lstrip("$"), 16))
            k += 1
        blocks.append((lo, hi, size, label, refs, bytes(data)))
        i = k
    return blocks


def looks_like_code(data, base):
    """True when the bytes decode as a run of instructions."""
    dec = 0
    off = 0
    while off < len(data):
        ent = nesdis.OPCODES.get(data[off])
        if ent is None:
            break
        length = nesdis.MODE_LEN[ent[1]]
        off += length
        dec += 1
    return dec


def analyse(lo, hi, size, label, refs, data, prg_base, prg_size):
    """The evidence for one block."""
    ev = []

    # 1. the consumers the tracer recorded
    if refs:
        consumers = ", ".join(f"{op} at ${at:04X}" for op, at in refs[:6])
        more = f" (+{len(refs) - 6} more)" if len(refs) > 6 else ""
        ev.append(f"read by {consumers}{more}")
    else:
        ev.append("no instruction in this dump references it")

    if not data:
        return ev, None

    # 2. is it a table of 16-bit pointers into the PRG?
    if len(data) % 2 == 0:
        ptrs = []
        ok = True
        for i in range(0, len(data), 2):
            v = data[i] | (data[i + 1] << 8)
            ptrs.append(v)
            if not (prg_base <= v < prg_base + prg_size):
                ok = False
        if ok and ptrs:
            ev.append(f"{len(ptrs)} 16-bit values, every one inside the PRG "
                      f"image - a pointer table")
            return ev, "pointer table"

    # 3. is it a table of small values?
    if len(data) >= 4 and max(data) < 0x40:
        c = Counter(data)
        top, n = c.most_common(1)[0]
        ev.append(f"values in $00-${max(data):02X}"
                  + (f", {n} of them are ${top:02X}" if n > 1 else ""))
        return ev, "byte table"

    # 4. padding?
    if len(set(data)) == 1:
        ev.append(f"every byte is ${data[0]:02X} - padding or fill")
        return ev, "padding"

    # 5. does it decode as code?
    dec = 0
    off = 0
    while off < len(data):
        ent = nesdis.OPCODES.get(data[off])
        if ent is None:
            break
        length = nesdis.MODE_LEN[ent[1]]
        off += length
        dec += 1
    if dec >= 4:
        ev.append(f"decodes as {dec} consecutive instructions - this is code "
                  "the tracer could not reach, not data")
        return ev, "code"

    ev.append(f"{len(data)} bytes, no single structure")
    return ev, None


def main():
    ap = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    ap.add_argument("project")
    ap.add_argument("--all", action="store_true",
                    help="list every block, not just the named ones")
    args = ap.parse_args()

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

    blocks = collect(d / "prg.s")
    scan_consumers(d / "prg.s", blocks)
    print(f"{d.name}: {len(blocks)} data blocks, PRG mapped at ${base:04X}\n")
    for lo, hi, size, label, refs, data in blocks:
        if not args.all and label is None:
            continue
        ev, kind = analyse(lo, hi, size, label, refs, data, base, len(prg))
        name = label or "(unnamed)"
        print(f"${lo:04X}-${hi:04X} ({size:5d} bytes)  {name}")
        for e in ev:
            print(f"      - {e}")
        print()


if __name__ == "__main__":
    main()