#!/usr/bin/env python3
"""Derive the Pass 3 zero-page alias set for Pac-Man (Tengen) (U).

The alias list is specific to that project; the script is kept with the
tooling because it records how the aliases were chosen.

Usage: python3 _tools/derive_zp_aliases.py <project dir>

Every alias here is read out of the Tengen dump's own code; nothing is
carried over from the sibling `Pac-Man (U)` dump.  Two kinds:

  global  the address has the same role at *every* reference, so the
          neutral ZpXX name is replaced everywhere.
  site    the address is reused for several roles; the alias applies only
          inside the named routines, where the role is proven.

Writes the alias block for locals.inc and a sites.tsv for
apply_zp_aliases.py, then prints what it did.
"""

import re
import sys
from pathlib import Path

PROJ = Path(sys.argv[1] if len(sys.argv) > 1 else
            "NROM/Pac-Man (Tengen) (U)")

# --- global aliases: address -> (name, why) --------------------------------
GLOBAL = {
    0x10: ("DispatchPtrLo", "low byte of the pointer the computed dispatch jumps through"),
    0x11: ("DispatchPtrHi", "high byte of that pointer"),
    0x12: ("ScoreLo", "low byte of the 16-bit score accumulator (AddScore/SubScore)"),
    0x13: ("ScoreHi", "high byte of the score accumulator"),
    0x14: ("ScoreAddendLo", "low byte of the value added to the score"),
    0x15: ("ScoreAddendHi", "high byte of the value added to the score"),
    0x3F: ("GameState", "index into the game-state dispatch table at $C253"),
    0x40: ("NmiFlag", "set to 1 by the game, waited on until the NMI clears it"),
    0x41: ("ScrollX", "written to PPUSCROLL"),
    0x42: ("ScrollY", "written to PPUSCROLL"),
    0x43: ("PpuCtrlShadow", "PPUCTRL shadow - always stored together with $2000"),
    0x45: ("OamDmaPage", "the page written to $4014 (OAMDMA)"),
    0x46: ("CocktailModeFlag", "ANDed with $47 to form the PPUCTRL nametable bit"),
    0x47: ("ActivePlayerFlag", "selects which controller the NMI reads"),
    0x48: ("GameplaySuspended", "$FF while the game is running; gates the level script"),
    0x4B: ("FrameCounter", "incremented every NMI; masked to pick the behaviour"),
    0x4D: ("Player1Input", "JOY1 shift register"),
    0x4E: ("Player2Input", "JOY2 shift register"),
    0x4F: ("ActiveInput", "the controller the NMI selected"),
    0x64: ("CopyrightVerified", "set to 1 once the signature is in RAM"),
    0x87: ("GhostDataIndex", "steps through the character table at $C66A"),
    0x88: ("GhostFlags", "mode flags for the ghost being updated"),
    0x89: ("TargetSlotIndex", "the pellet slot picked as the ghost's target"),
    0xB7: ("FruitIndex", "the fruit type counter"),
    0xB8: ("GhostSlotIndex", "the ghost slot, indexed by X"),
    0xB9: ("GhostState", "per-ghost state, indexed by X"),
    0xCD: ("PositionDeltaX", "the computed X delta"),
    0xCE: ("PositionDeltaY", "the computed Y delta"),
    0xE3: ("AnimationTimer", "the animation frame timer"),
    0xE4: ("AnimationFrameIndex", "the animation frame index"),
}

# --- site aliases: (routine start, routine end, address, name, why) --------
SITES = [
    (0xC27D, 0xC2DC, 0x00, "PtrLo", "low byte of the VRAM address DrawMazeGraphics writes to"),
    (0xC27D, 0xC2DC, 0x01, "PtrHi", "high byte of that VRAM address"),
    (0xC589, 0xC5B0, 0x00, "PtrLo", "low byte of the character-data pointer"),
    (0xC589, 0xC5B0, 0x01, "PtrHi", "high byte of that pointer"),
    (0xDAB0, 0xDBAC, 0x00, "PtrLo", "low byte of the source pointer UpdateMazeTiles reads"),
    (0xDAB0, 0xDBAC, 0x01, "PtrHi", "high byte of that pointer"),
    (0xDAB0, 0xDBAC, 0x02, "Ptr2Lo", "low byte of the destination pointer"),
    (0xDAB0, 0xDBAC, 0x03, "Ptr2Hi", "high byte of that pointer"),
    (0xE9D8, 0xEA38, 0x00, "PtrLo", "low byte of the maze-data pointer"),
    (0xE9D8, 0xEA38, 0x01, "PtrHi", "high byte of that pointer"),
    (0xE9D8, 0xEA38, 0x02, "Ptr2Lo", "low byte of the cell-state pointer"),
    (0xE9D8, 0xEA38, 0x03, "Ptr2Hi", "high byte of that pointer"),
    (0xC8ED, 0xC976, 0x00, "PtrLo", "low byte of the 32-bit ghost pointer"),
    (0xC8ED, 0xC976, 0x01, "PtrHi", "byte 1 of that pointer"),
    (0xC8ED, 0xC976, 0x02, "Ptr2Lo", "byte 2 of that pointer"),
    (0xC8ED, 0xC976, 0x03, "Ptr2Hi", "byte 3 of that pointer"),
]

CPU_RE = re.compile(r"; CPU:\$([0-9A-F]{4}) PRG:")

# Numeric operands that carry an alias but are not plain ZpXX tokens, so
# apply_symbols.py cannot rewrite them on its own:
#   (CPU address, old token, new name)
EXTRA_SITES = [
    (0xC048, "a:$0000", "Zp00"),   # the zero-page clear loop's base address
    (0xC089, "a:$0000", "Zp00"),   # the full zero-page clear
    (0xC250, "($0010)", "(DispatchPtrLo)"),
    (0xD3CA, "($0010)", "(DispatchPtrLo)"),
    (0xD90F, "($0010)", "(DispatchPtrLo)"),
    (0xD9B6, "($0010)", "(DispatchPtrLo)"),
]


def main():
    prg = PROJ / "prg.s"
    lines = prg.read_text().split("\n")

    # neutral name for each address, straight out of locals.inc
    loc = (PROJ / "locals.inc").read_text()
    neutral = {}
    for m in re.finditer(r"^(Zp[0-9A-F]{2})\s*=\s*\$([0-9A-F]{2})", loc, re.M):
        neutral[int(m.group(2), 16)] = m.group(1)

    # every reference site, by CPU address.  Both the neutral name and an
    # already-applied alias are recognised, so the script is idempotent.
    refs = []
    for i, line in enumerate(lines, 1):
        m = CPU_RE.search(line)
        if not m:
            continue
        a = int(m.group(1), 16)
        for n in re.findall(r"\b(?:Zp[0-9A-F]{2}|PtrLo|PtrHi|Ptr2Lo|Ptr2Hi)\b", line):
            refs.append((a, i, n))

    alias_of = {}
    for a, (name, _why) in GLOBAL.items():
        alias_of[a] = name

    site_rows = []
    for lo, hi, a, name, _why in SITES:
        for addr, lineno, n in refs:
            # Accept the neutral name or the alias itself, so the script is
            # idempotent and can be re-run after the aliases are applied.
            if lo <= addr <= hi and n in (neutral.get(a), name):
                site_rows.append((lineno, name))

    # --- write the alias block for locals.inc -----------------------------
    # Strip any block a previous run left, so the script is idempotent.
    marker = "; Semantic zero-page aliases (Pass 3)"
    if marker in loc:
        loc = loc[:loc.index(marker)].rstrip("\n") + "\n"

    block = ["", "; ---------------------------------------------------------------------------",
             marker,
             ";",
             "; Every alias below was read out of this dump's own code.  An address whose",
             "; role is the same at every reference is renamed everywhere; an address that",
             "; is reused for several roles keeps the neutral ZpXX name outside the",
             "; routines where the role is proven, and the per-site choice is recorded in",
             "; sites.tsv.",
             "; ---------------------------------------------------------------------------"]
    for a in sorted(alias_of):
        name = alias_of[a]
        why = GLOBAL.get(a, (None, ""))[1]
        block.append(f"{name:<24}= ${a:02X} ; semantic alias for {neutral.get(a, '?')}"
                     + (f" ({why})" if why else ""))

    # The per-site aliases also need a definition; their role differs from
    # site to site, so the reason records the routine they belong to.
    site_alias = {}
    for lo, hi, a, name, why in SITES:
        site_alias.setdefault((a, name), (why, lo, hi))
    for (a, name), (why, lo, hi) in sorted(site_alias.items()):
        block.append(f"{name:<24}= ${a:02X} ; semantic alias for {neutral.get(a, '?')} "
                     f"({why}; applied only inside ${lo:04X}-${hi:04X}, elsewhere the "
                     f"neutral name is kept)")
    (PROJ / "locals.inc").write_text(loc.rstrip("\n") + "\n" + "\n".join(block) + "\n")

    # --- write sites.tsv ---------------------------------------------------
    seen = {}
    for lineno, name in site_rows:
        seen[lineno] = name
    out = ["# Pac-Man (Tengen) (U) - per-site zero-page alias choice.",
           "# Format: <line in prg.s><TAB><old token><TAB><new name>",
           "# An address that carries several roles keeps the neutral ZpXX name at",
           "# the sites where the other role applies.  The last six rows resolve the",
           "# numeric operands that apply_symbols.py cannot rewrite on its own.",
           ""]
    for lineno in sorted(seen):
        name = seen[lineno]
        old = next(n for a, ln, n in refs if ln == lineno)
        out.append(f"{lineno}\t{old}\t{name}")
    # Numeric operands that carry an alias but are not plain ZpXX tokens.
    for cpu, old, new in EXTRA_SITES:
        lineno = next((ln for a, ln, _n in refs if a == cpu), None)
        if lineno is None:
            for ln, line in enumerate(lines, 1):
                if (m := CPU_RE.search(line)) and int(m.group(1), 16) == cpu:
                    lineno = ln
                    break
        if lineno is None:
            raise SystemExit(f"error: no source line for ${cpu:04X}")
        out.append(f"{lineno}\t{old}\t{new}")
    (PROJ / "sites.tsv").write_text("\n".join(out) + "\n")

    print(f"{len(alias_of)} aliases, {len(seen)} per-site choices")


if __name__ == "__main__":
    main()