#!/usr/bin/env python3
"""Evidence-graded function-boundary inventory for the USA executable.

This replaces address-gap guessing with a deterministic scan that labels every
candidate function start with the evidence that supports it:

  entry     the PS-X EXE header's entry PC.
  jal       the target of a `jal` instruction inside the payload. A direct call
            target is a function entry; this is the strongest grade.
  prologue  an address whose first instruction is `addiu sp,sp,-N` (N > 0) and
            whose second is a `sw ...,off(sp)`. This is the standard non-leaf
            prologue; it can be a false positive if such a sequence appears
            inside another function.
  ghidra    an address present in a Ghidra function list supplied by the caller
            (`--ghidra FILE`, one hex address per line). Ghidra's auto-analysis
            both misses functions and invents them, so this grade is neither
            necessary nor sufficient.

Output is a sorted TSV of `address<TAB>grades` with a short header. It contains
addresses only -- never instruction bytes -- and the input executable is only
read from a caller-supplied path. Nothing here promotes a candidate to a match:
a match still requires an instruction-identical comparison and the full-binary
gate.

The library-versus-game-code question is **unresolved**; no grade distinguishes
a PsyQ library routine from game code.

Exit codes: 0 success, 2 usage or environment error.
"""

from __future__ import annotations

import argparse
from pathlib import Path
import struct
import sys
from typing import Sequence


EXE_MAGIC = b"PS-X EXE"
HEADER_SIZE = 0x800
PAYLOAD_LMA = 0x800

GRADE_ORDER = ("entry", "jal", "prologue", "ghidra")


class ToolError(Exception):
    """A usage or environment problem; maps to exit code 2."""


def parse_hex(text: str, label: str) -> int:
    try:
        return int(text, 16)
    except ValueError as exc:
        raise ToolError(f"{label}: not a hex address: {text!r}") from exc


def require_file(path: Path, label: str) -> Path:
    if not path.is_file():
        raise ToolError(f"{label} is not a regular file: {path}")
    return path


def fresh_path(path: Path, label: str) -> Path:
    if path.exists() or path.is_symlink():
        raise ToolError(f"{label} already exists: {path}")
    return path


def parse_psx_exe(header: bytes) -> tuple[int, int, int]:
    if len(header) < HEADER_SIZE:
        raise ToolError("executable is smaller than a PS-X EXE header")
    if header[:8] != EXE_MAGIC:
        raise ToolError("executable does not carry the PS-X EXE magic")
    entry, _gp, text_address, text_size = struct.unpack_from("<IIII", header, 0x10)
    if text_size == 0:
        raise ToolError("PS-X EXE header declares an empty payload")
    return entry, text_address, text_size


def load_ghidra(path: Path | None) -> set[int]:
    if path is None:
        return set()
    addresses: set[int] = set()
    for number, raw in enumerate(path.read_text(encoding="utf-8").splitlines(), 1):
        line = raw.split("#", 1)[0].strip()
        if not line:
            continue
        addresses.add(parse_hex(line, f"ghidra line {number}"))
    return addresses


def scan(payload: bytes, text_address: int, entry: int, ghidra: set[int]) -> dict[int, set[str]]:
    """Return address -> evidence grades for every candidate start."""
    grades: dict[int, set[str]] = {}
    end = text_address + len(payload)

    def add(address: int, grade: str) -> None:
        if text_address <= address < end and address % 4 == 0:
            grades.setdefault(address, set()).add(grade)

    add(entry, "entry")

    for address in ghidra:
        add(address, "ghidra")

    words = len(payload) // 4
    for index in range(words):
        word = struct.unpack_from("<I", payload, index * 4)[0]
        pc = text_address + index * 4
        opcode = word >> 26
        if opcode == 0x03:  # jal
            target = ((pc + 4) & 0xF0000000) | ((word & 0x03FFFFFF) << 2)
            add(target, "jal")
        elif opcode == 0x09:  # addiu sp,sp,-N  followed by  sw ...,off(sp)
            rs = (word >> 21) & 0x1F
            rt = (word >> 16) & 0x1F
            immediate = word & 0xFFFF
            signed = immediate - 0x10000 if immediate & 0x8000 else immediate
            if rs == 29 and rt == 29 and signed < 0 and index + 1 < words:
                following = struct.unpack_from("<I", payload, (index + 1) * 4)[0]
                if (following >> 26) == 0x2B and ((following >> 21) & 0x1F) == 29:
                    add(pc, "prologue")

    return grades


def format_inventory(grades: dict[int, set[str]]) -> str:
    lines = [
        "# Syphon Filter 3 (USA) function-boundary inventory.",
        "# Columns: address<TAB>comma-separated evidence grades.",
        "# Grades: entry, jal, prologue, ghidra. Addresses only; no bytes.",
        "# Library-versus-game-code is unresolved; no grade answers it.",
    ]
    for address in sorted(grades):
        names = [grade for grade in GRADE_ORDER if grade in grades[address]]
        lines.append(f"0x{address:08X}\t{','.join(names)}")
    return "\n".join(lines) + "\n"


def command_scan(args: argparse.Namespace) -> int:
    exe_path = require_file(args.exe, "executable")
    out = fresh_path(args.out, "output")
    data = exe_path.read_bytes()
    entry, text_address, text_size = parse_psx_exe(data)
    payload = data[PAYLOAD_LMA:PAYLOAD_LMA + text_size]
    ghidra = load_ghidra(args.ghidra)
    grades = scan(payload, text_address, entry, ghidra)

    counts = {grade: 0 for grade in GRADE_ORDER}
    for names in grades.values():
        for grade in names:
            counts[grade] += 1
    out.parent.mkdir(parents=True, exist_ok=True)
    out.write_text(format_inventory(grades), encoding="ascii")

    print(f"candidates={len(grades)}")
    for grade in GRADE_ORDER:
        print(f"grade_{grade}={counts[grade]}")
    print(f"output={out}")
    return 0


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter
    )
    subparsers = parser.add_subparsers(dest="command", required=True)
    scan_parser = subparsers.add_parser("scan", help="scan an executable")
    scan_parser.add_argument("--exe", required=True, type=Path)
    scan_parser.add_argument("--ghidra", type=Path, default=None,
                             help="Ghidra function list, one hex address per line")
    scan_parser.add_argument("--out", required=True, type=Path)
    scan_parser.set_defaults(handler=command_scan)
    return parser


def main(argv: Sequence[str] | None = None) -> int:
    parser = build_parser()
    args = parser.parse_args(argv)
    try:
        return args.handler(args)
    except ToolError as exc:
        print(f"error: {exc}", file=sys.stderr)
        return 2
    except OSError as exc:
        print(f"error: {exc}", file=sys.stderr)
        return 2


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