#!/usr/bin/env python3
"""
Generate a byte-exact assembly stub of the ROM code section and a
function-boundary map.

Outputs:
  - asm/all_stubs.s            byte-exact code-section stub
  - build/stub_map.json        name -> {addr, size} for each function
"""

import json
from pathlib import Path

ROM_PATH = Path("bit Generations - Orbital (Japan) (En).gba")
FUNC_LIST = Path("function_list_v2.txt")
ALL_FILE = Path("asm/all_stubs.s")
MAP_FILE = Path("build/stub_map.json")

ROM_BASE = 0x08000000
CODE_START = 0x080000C0
CODE_END = 0x0801F30C
CODE_SIZE = CODE_END - CODE_START  # 127564

def parse_functions(path):
    rows = []
    with open(path) as f:
        for line in f:
            parts = line.strip().split(None, 2)
            if len(parts) >= 2:
                addr = int(parts[0], 16)
                size = int(parts[1])
                name = parts[2] if len(parts) > 2 else f"sub_{addr:X}"
                rows.append((addr, size, name))
    rows.sort()
    # deduplicate by address (keep first)
    seen_addr = set()
    out = []
    for a, s, n in rows:
        if a not in seen_addr:
            seen_addr.add(a)
            out.append((a, s, n))
    # disambiguate duplicate names (same name at different address)
    used = {}
    final = []
    for a, s, n in out:
        if n in used:
            n = f"{n}_{a:08X}"
        else:
            used[n] = a
        final.append((a, s, n))
    return final

def main():
    rom = ROM_PATH.read_bytes()
    funcs = parse_functions(FUNC_LIST)
    code = rom[CODE_START - ROM_BASE:CODE_END - ROM_BASE]

    MAP_FILE.parent.mkdir(parents=True, exist_ok=True)
    ALL_FILE.parent.mkdir(parents=True, exist_ok=True)

    # Build a map: code_offset -> (name, size)
    func_map = {}
    for addr, size, name in funcs:
        off = addr - CODE_START
        if 0 <= off < CODE_SIZE and size >= 2:
            func_map[off] = (name, size)

    with open(ALL_FILE, "w") as af:
        af.write("\t.thumb\n\t.code 16\n\n")

        i = 0
        while i < CODE_SIZE:
            if i in func_map:
                name, size = func_map[i]
                af.write(f"\t.global {name}\n\t.thumb_func\n{name}:\n")
                end = min(i + size, CODE_SIZE)
                while i < end - 1:
                    hw = int.from_bytes(code[i:i+2], 'little')
                    af.write(f"\t.short 0x{hw:04X}\n")
                    i += 2
                if i < end:
                    af.write(f"\t.byte 0x{code[i]:02X}\n")
                    i += 1
            else:
                # raw code bytes (pools / data / alignment / function gaps)
                af.write(f"\t.short 0x{int.from_bytes(code[i:i+2], 'little'):04X}\n")
                i += 2

    with open(MAP_FILE, "w") as f:
        rows = [{"name": n, "addr": a, "size": s} for a, s, n in funcs
                if CODE_START <= a < CODE_END and s >= 2]
        # Size-0 ALIAS_<name> lines: pure symbols for alternate entry
        # points (shared-code idiom), no byte ownership.
        for a, s, n in funcs:
            if s == 0 and n.startswith("ALIAS_") and CODE_START <= a < CODE_END:
                rows.append({"name": n[len("ALIAS_"):], "addr": a,
                             "size": 0, "alias": True})
        json.dump(rows, f, indent=2)

    print(f"Wrote {ALL_FILE} ({CODE_SIZE} code bytes, {len(func_map)} named functions)")
    print(f"Wrote {MAP_FILE}")

if __name__ == "__main__":
    main()