#!/usr/bin/env python3
"""
Function segmentation analyzer for bit Generations: Orbital.

Builds a ground-truth function list from the ROM + address list, using real
return instructions (bx lr, pop {..,pc}) and padding detection to estimate
true function sizes rather than relying on raw address-gap arithmetic.

Outputs:
  - build/functions.json  (programmatic ground-truth list)
  - build/functions.txt   (human-readable table)
  - build/segmentation_stats.txt
"""
import json
import os
import re
from dataclasses import asdict, dataclass
from pathlib import Path

ROM_PATH = Path("bit Generations - Orbital (Japan) (En).gba")
FUNC_LIST = Path("function_list_v2.txt")
OUT_DIR = Path("build")

ROM_BASE = 0x08000000
CODE_START = 0x080000C0
CODE_END = 0x0801F30C

@dataclass
class FuncEntry:
    addr: int
    listed_name: str
    listed_size: int
    real_size: int
    last_return_offset: int
    return_type: str
    has_bl: bool
    has_pc_ldr: bool
    reloc_words: list  # offsets within function of relocatable words
    thumb_instructions: int
    padding_bytes: int
    status: str  # "clean", "reloc-bl", "reloc-pool", "reloc-both", "no-ret", "pad-only"

def read_word_le(data: bytes, offset: int) -> int:
    return int.from_bytes(data[offset:offset+2], 'little')

def is_32bit_thumb(hw1: int) -> bool:
    return (hw1 >> 11) in (0b11101, 0b11110, 0b11111)

def detect_returns(func_bytes: bytes):
    """Walk the function and record return-type instructions."""
    returns = []
    i = 0
    n = len(func_bytes)
    while i < n - 1:
        hw = read_word_le(func_bytes, i)
        if hw == 0x4770:
            returns.append((i, "bx lr"))
        elif (hw >> 8) == 0xBD and (hw & 0x100):
            returns.append((i, "pop {..,pc}"))
        elif (hw & 0xFF87) == 0x4700:
            returns.append((i, "bx rN"))
        # 32-bit encoding
        if is_32bit_thumb(hw):
            i += 4
        else:
            i += 2
    return returns

def detect_relocations(func_bytes: bytes, func_addr: int):
    """Detect BL immediates and PC-relative LDR pool references."""
    relocs = []
    has_bl = False
    has_pool_ldr = False
    i = 0
    n = len(func_bytes)
    while i < n - 3:
        hw1 = read_word_le(func_bytes, i)
        hw2 = read_word_le(func_bytes, i + 2)
        if (hw1 >> 11) == 0b11110 and (hw2 >> 11) == 0b11111:
            # BL: mark the 4-byte immediate for masking
            relocs.append(("bl", i))
            has_bl = True
            i += 4
            continue
        i += 2
    for j in range(0, n - 1, 2):
        hw = read_word_le(func_bytes, j)
        if (hw >> 11) == 0b01001:
            has_pool_ldr = True
    return relocs, has_bl, has_pool_ldr

def parse_address_list(path: Path):
    entries = []
    with open(path) as f:
        for line in f:
            m = re.search(r"([0-9A-Fa-f]{6,8})\s+(\d+)\s+(.+)", line)
            if m:
                entries.append((int(m.group(1), 16), int(m.group(2)), m.group(3).strip()))
    return sorted(entries)

def classify_gap(gap: bytes):
    """Classify inter-function bytes as literal-pool words / padding / data."""
    pools = 0
    i = 0
    while i + 3 < len(gap):
        w = int.from_bytes(gap[i:i + 4], "little")
        if (w >> 24) in (0x08, 0x02, 0x03, 0x04) and (w & 0x3) == 0:
            # plausible only if it points at code/data regions used by neighbors;
            # conservatively count aligned plausible pointers
            pools += 4
            i += 4
        else:
            break
    rest = gap[i:]
    pad = 0
    while pad < len(rest) and rest[len(rest) - 1 - pad] == 0:
        pad += 1
    return {"pool_bytes": pools, "pad_bytes": pad, "other_bytes": len(rest) - pad,
            "total": len(gap)}

def main():
    rom = ROM_PATH.read_bytes()
    addr_list = parse_address_list(FUNC_LIST)
    os.makedirs(OUT_DIR, exist_ok=True)

    results = []
    addrs = sorted(a for a, _, _ in addr_list)
    nxt = {a: (addrs[i + 1] if i + 1 < len(addrs) else CODE_END)
           for i, a in enumerate(addrs)}
    for addr, listed_size, name in addr_list:
        if not (CODE_START <= addr < CODE_END):
            continue
        if listed_size < 2:
            continue
        offset = addr - ROM_BASE
        hi = min(nxt[addr] - ROM_BASE, offset + listed_size + 256, len(rom))
        span = rom[offset:hi]
        returns = detect_returns(span)
        if not returns:
            last_ret_off = -1
            ret_type = "none"
            real_size = min(listed_size, hi - offset)
        else:
            last_ret_off = returns[-1][0] + 2
            ret_type = returns[-1][1]
            # real body ends at the LAST return inside this function's span.
            # mid-function bx lr / pop are early exits, not ends.
            real_size = last_ret_off
        body = rom[offset:offset + real_size]
        relocs, has_bl, has_pool_ldr = detect_relocations(body, addr)
        gap = rom[offset + real_size:hi]
        gc = classify_gap(gap)
        p = gc["pad_bytes"]

        # status classification
        if not returns:
            status = "no-ret"
        elif has_bl and has_pool_ldr:
            status = "reloc-both"
        elif has_bl:
            status = "reloc-bl"
        elif has_pool_ldr:
            status = "reloc-pool"
        elif p == real_size and real_size <= 2:
            status = "pad-only"
        else:
            status = "clean"

        results.append(FuncEntry(
            addr=addr,
            listed_name=name,
            listed_size=listed_size,
            real_size=real_size,
            last_return_offset=last_ret_off,
            return_type=ret_type,
            has_bl=has_bl,
            has_pc_ldr=has_pool_ldr,
            reloc_words=[r[1] for r in relocs],
            thumb_instructions=max(1, (real_size + 1) // 2),
            padding_bytes=p,
            status=status,
        ))

    # write outputs
    with open(OUT_DIR / "functions.json", "w") as f:
        json.dump([asdict(r) for r in results], f, indent=2)
    with open(OUT_DIR / "functions.txt", "w") as f:
        f.write(f"{'addr':>10s}  {'listed':>6s}  {'real':>5s}  {'ret':>12s}  {'pad':>3s}  {'status':>12s}  name\n")
        for r in results:
            f.write(f"0x{r.addr:08X}  {r.listed_size:6d}  {r.real_size:5d}  {r.return_type:>12s}  {r.padding_bytes:3d}  {r.status:>12s}  {r.listed_name}\n")

    # summary
    total_listed = sum(r.listed_size for r in results)
    total_real = sum(r.real_size for r in results)
    total_pad = sum(r.padding_bytes for r in results)
    clean = sum(r.real_size for r in results if r.status == "clean")
    by_status = {}
    for r in results:
        by_status.setdefault(r.status, 0)
        by_status[r.status] += r.real_size

    with open(OUT_DIR / "segmentation_stats.txt", "w") as f:
        f.write(f"Functions analyzed: {len(results)}\n")
        f.write(f"Sum listed sizes: {total_listed}\n")
        f.write(f"Sum real sizes (pad-stripped): {total_real}\n")
        f.write(f"Padding stripped: {total_pad}\n")
        f.write(f"Code section size: {CODE_END - CODE_START}\n")
        f.write(f"Clean reloc-free functions: {clean} bytes ({100*clean/127564:.1f}% of code section)\n")
        f.write("\nBy status:\n")
        for k in ("clean", "reloc-bl", "reloc-pool", "reloc-both", "no-ret", "pad-only"):
            v = by_status.get(k, 0)
            f.write(f"  {k:>12s}: {v:7d} bytes ({100*v/127564:5.1f}%)\n")
        f.write("\nTop 20 largest real functions:\n")
        for r in sorted(results, key=lambda x: -x.real_size)[:20]:
            f.write(f"  0x{r.addr:08X} {r.real_size:6d}B  {r.listed_name}\n")

    print(f"Wrote {OUT_DIR}/functions.json ({len(results)} entries)")
    print(f"Real code bytes (pad-stripped): {total_real}")
    print(f"Clean reloc-free: {clean} bytes ({100*clean/127564:.1f}%)")

if __name__ == "__main__":
    main()