#!/usr/bin/env python3
"""Extract native GodotSpikeBridge sound registrations from the game binary.

This is a static AArch64 pattern extractor for the registration table
identified in Ghidra as FUN_005d6b90 -> FUN_006ec060(name, id, -1).
It does not execute the target binary.
"""

from __future__ import annotations

import argparse
import csv
import struct
from pathlib import Path


GHIDRA_BASE_DELTA = 0x100000
GHIDRA_REGISTER_SOUND = 0x006EC060
TEXT_START = 0x00460000
TEXT_END = 0x0258A7B0


def u32(data: bytes, offset: int) -> int:
    return struct.unpack_from("<I", data, offset)[0]


def decode_bl_target(pc: int, insn: int) -> int | None:
    if insn & 0xFC000000 != 0x94000000:
        return None
    imm = insn & 0x03FFFFFF
    if imm & 0x02000000:
        imm -= 0x04000000
    return pc + imm * 4


def decode_movz_w2(insn: int) -> int | None:
    if insn & 0x7F80001F != 0x52800002:
        return None
    shift = (insn >> 21) & 0x3
    if shift != 0:
        return None
    return (insn >> 5) & 0xFFFF


def decode_adrp(pc: int, insn: int) -> tuple[int, int] | None:
    if insn & 0x9F000000 != 0x90000000:
        return None
    reg = insn & 0x1F
    immlo = (insn >> 29) & 0x3
    immhi = (insn >> 5) & 0x7FFFF
    imm = (immhi << 2) | immlo
    if imm & (1 << 20):
        imm -= 1 << 21
    return reg, (pc & ~0xFFF) + imm * 0x1000


def decode_add_imm(insn: int) -> tuple[int, int, int] | None:
    if insn & 0xFFC00000 != 0x91000000:
        return None
    rd = insn & 0x1F
    rn = (insn >> 5) & 0x1F
    shift = (insn >> 22) & 0x3
    imm = (insn >> 10) & 0xFFF
    if shift == 1:
        imm <<= 12
    elif shift != 0:
        return None
    return rd, rn, imm


def read_c_string(data: bytes, offset: int) -> str:
    end = data.find(b"\0", offset)
    if end < 0:
        end = min(len(data), offset + 256)
    return data[offset:end].decode("ascii", errors="replace")


def extract(data: bytes, target_file_addr: int) -> list[dict[str, object]]:
    rows: list[dict[str, object]] = []
    start = min(TEXT_START, len(data))
    end = min(TEXT_END, len(data))
    for pc in range(start, end, 4):
        insn = u32(data, pc)
        if decode_bl_target(pc, insn) != target_file_addr:
            continue
        sound_id = None
        string_addr = None
        for cursor in range(pc - 4, max(TEXT_START, pc - 0x100), -4):
            prior = u32(data, cursor)
            if sound_id is None:
                sound_id = decode_movz_w2(prior)
            add_info = decode_add_imm(prior)
            if add_info is not None and cursor >= TEXT_START + 4:
                rd, rn, add_imm = add_info
                adrp_info = decode_adrp(cursor - 4, u32(data, cursor - 4))
                if adrp_info is not None:
                    adrp_reg, adrp_addr = adrp_info
                    if rd == rn and rn == adrp_reg:
                        string_addr = adrp_addr + add_imm
            if sound_id is not None and string_addr is not None:
                break
        rows.append(
            {
                "name": read_c_string(data, string_addr) if string_addr is not None else "",
                "native_sound_id": sound_id if sound_id is not None else "",
                "register_call_ghidra": f"0x{pc + GHIDRA_BASE_DELTA:08x}",
                "register_call_file_offset": f"0x{pc:08x}",
                "string_ghidra": f"0x{string_addr + GHIDRA_BASE_DELTA:08x}" if string_addr is not None else "",
                "string_file_offset": f"0x{string_addr:08x}" if string_addr is not None else "",
                "register_function_ghidra": f"0x{GHIDRA_REGISTER_SOUND:08x}",
            }
        )
    return rows


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--game-binary", type=Path, default=Path("games/pokemon_pro/game"))
    parser.add_argument("--out", type=Path, default=Path("analysis/radium-image/native_sound_id_map.tsv"))
    args = parser.parse_args()

    data = args.game_binary.read_bytes()
    rows = extract(data, GHIDRA_REGISTER_SOUND - GHIDRA_BASE_DELTA)
    args.out.parent.mkdir(parents=True, exist_ok=True)
    fields = [
        "name",
        "native_sound_id",
        "register_call_ghidra",
        "register_call_file_offset",
        "string_ghidra",
        "string_file_offset",
        "register_function_ghidra",
    ]
    with args.out.open("w", newline="") as f:
        writer = csv.DictWriter(f, fieldnames=fields, delimiter="\t")
        writer.writeheader()
        writer.writerows(rows)
    print(f"wrote {args.out} rows={len(rows)}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())