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

Phase 6 produced evidence-graded function *starts* (`config/function_inventory.tsv`)
and no evidence at all for where a function *ends*: every region end in
`config/regions.tsv` was derived by hand. This tool derives the end from the
payload's own control flow and states how much the result can be trusted.

## The model

A **hard start** is a candidate whose evidence is a direct call or the header
entry (`jal` or `entry` grade). A **soft start** is a `prologue`/`ghidra`-only
candidate. Phase 7 measurement showed soft starts are frequently *false
positives inside a real function body* -- e.g. a body's second instruction can
itself look like a prologue -- so soft starts are **not** used as walk
boundaries, and the walk is never stopped by them.

From each hard start the tool walks control flow (linear fall-through, branch
and jump targets, delay slots, calls falling through to the return address) and
records the highest reachable instruction. The walk stops at:

  * a return (`jr ra`) -- the delay slot is included;
  * a register jump (`jr $rs`, rs != ra), i.e. a switch/jump table, whose
    targets cannot be followed statically;
  * `break`, which traps;
  * a tail `j` whose target is outside this function (behind its start, or a
    known start);
  * another **hard** start, which is a boundary the walk may not cross.

`syscall` is deliberately **not** a terminator: on this target it is the BIOS
call instruction and it returns to the instruction after it. Treating it as a
return truncated the two 16-byte BIOS stubs at `0x80103FCC` and `0x80103FEC` to
8 bytes, which is unreachable from C because `cc1` always appends an epilogue
(found by worker C in Phase 8; the stub's four `jal` callers and the absence of
any reference to the following address are the confirming evidence).

The extent end is the first byte after the last reachable instruction. For a
call this is the return address; for a return, the instruction after the delay
slot.

## Grades

  exact        the walk terminated at a return / `break` / tail jump and stayed
               inside its own region. This is the grade a match candidate wants:
               the end is the function's own last byte + 1. All 115 regions
               registered in Phases 6-8 reproduce at this grade.
  fallthrough  no terminal was reachable; the walk ran into the next hard
               start. The extent ends there. Typically a function that ends in
               a call to something that does not return.
  indirect     the walk hit a register jump (`jr $rs`), so its result is a
               **lower bound**: the switch's case bodies are not reachable
               statically. The recorded end is the next hard start, which is
               the best available bound and is *not* verified.
  escape       the walk crossed the next hard start, so control left this
               function's region (a shared tail block, an interleaved body, or
               bytes decoded as code). The extent is clipped to the next hard
               start and the crossing instruction is recorded as evidence.
  outside      the walk left the declared payload.
  runaway      the walk exceeded the size cap without terminating.
  contained    a **soft** candidate lying inside an `exact` extent. Strong
               evidence that it is not an independent function.
  standalone   a soft candidate not covered by any `exact` extent. Unresolved:
               it may be a function reached only indirectly, or a false
               positive in data.

`contained` and `standalone` candidates carry no extent: their `end`, `size`,
`next` and `gap` fields are `-`.

## Output

A sorted TSV of `address<TAB>end<TAB>size<TAB>next<TAB>gap<TAB>grade<TAB>evidence`:

  * `next` is the next hard start (the next row that has an extent);
  * `gap` is `next - end` -- 0 means the function is packed against the next
    one, which is what a registerable region needs; a positive gap is padding
    (alignment) that must stay out of the region;
  * `evidence` is a `;`-separated list of `key=value` tokens.

It contains **addresses, sizes and grade names only -- never instruction bytes**.
The executable is read only from the caller-supplied path.

Nothing here promotes a candidate to a match: a match still requires an
instruction-identical `sf3_match range` comparison and the clean full-binary
gate.

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

from __future__ import annotations

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


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

HARD_GRADES = frozenset({"jal", "entry"})

# Grade order for reporting.
GRADES = (
    "exact",
    "fallthrough",
    "indirect",
    "escape",
    "outside",
    "runaway",
    "contained",
    "standalone",
)

HEADER_LINES = (
    "# Syphon Filter 3 (USA) function extents.",
    "# Columns: address<TAB>end<TAB>size<TAB>next<TAB>gap<TAB>grade<TAB>evidence.",
    "# Grades: exact, fallthrough, indirect, escape, outside, runaway,",
    "#         contained, standalone. Addresses and sizes only; no bytes.",
    "# A dash means the field does not apply (a non-function candidate).",
    "# 'exact' is the only grade a match candidate should rely on;",
    "# 'contained' is evidence a soft candidate is NOT a function start.",
    "# Regenerate: ./tools/sf3_extents scan --exe '<exe>' \\",
    "#   --inventory config/function_inventory.tsv --out config/function_extents.tsv --force",
)


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


class Disagreement(Exception):
    """A verification failure; maps to exit code 1."""


# --------------------------------------------------------------------------
# parsing helpers
# --------------------------------------------------------------------------


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 resolve_output(path: Path, force: bool) -> Path:
    if path.exists() or path.is_symlink():
        if not force:
            raise ToolError(f"output already exists (use --force to overwrite): {path}")
        if not path.is_file() or path.is_symlink():
            raise ToolError(f"output is not a regular file: {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_inventory(path: Path) -> dict[int, set[str]]:
    """Read `address<TAB>grades` rows, skipping comments and blank lines."""
    grades: dict[int, set[str]] = {}
    for number, raw in enumerate(path.read_text(encoding="utf-8").splitlines(), 1):
        line = raw.split("#", 1)[0].strip()
        if not line:
            continue
        fields = line.split("\t")
        if len(fields) != 2:
            raise ToolError(f"inventory line {number}: expected 'address<TAB>grades'")
        address = parse_hex(fields[0].strip(), f"inventory line {number}")
        names = {name.strip() for name in fields[1].split(",") if name.strip()}
        if not names:
            raise ToolError(f"inventory line {number}: no grades")
        if address in grades:
            raise ToolError(f"inventory line {number}: duplicate address {fields[0]!r}")
        grades[address] = names
    if not grades:
        raise ToolError("inventory is empty")
    return grades


def load_regions(path: Path) -> list[tuple[int, int, str]]:
    """Read `start<TAB>end<TAB>source[<TAB>overrides]` rows."""
    regions: list[tuple[int, int, str]] = []
    for number, raw in enumerate(path.read_text(encoding="utf-8").splitlines(), 1):
        line = raw.split("#", 1)[0].strip()
        if not line:
            continue
        fields = line.split("\t")
        if len(fields) < 3:
            raise ToolError(f"region line {number}: expected at least three fields")
        start = parse_hex(fields[0].strip(), f"region line {number}")
        end = parse_hex(fields[1].strip(), f"region line {number}")
        if end <= start:
            raise ToolError(f"region line {number}: end is not after start")
        regions.append((start, end, fields[2].strip()))
    return regions


# --------------------------------------------------------------------------
# instruction decoding and the control-flow walk
# --------------------------------------------------------------------------


class Walk:
    """The result of walking one candidate start.

    `reach` is the highest reachable instruction address + 4: the whole
    reachable body is explored, not just the first path that terminates, so a
    function whose paths return at different addresses still reports its true
    extent. `boundary` is the lowest **hard** start the walk ran into, which is
    a boundary it may not cross. `escape_at` records the first branch or jump
    whose target left the region, which is what distinguishes a function that
    *falls through* into its neighbour from one that *jumps over* it.
    """

    __slots__ = ("reach", "boundary", "indirect_at", "terminals", "escape_at", "overflow")

    def __init__(self, reach: int, boundary: int | None, indirect_at: int | None,
                 terminals: list[tuple[str, int, int | None]], escape_at: tuple[int, int, str] | None,
                 overflow: tuple[str, int] | None) -> None:
        self.reach = reach
        self.boundary = boundary
        self.indirect_at = indirect_at
        self.terminals = terminals
        self.escape_at = escape_at
        self.overflow = overflow

    def terminal_evidence(self) -> str:
        if not self.terminals:
            return "term=none"
        kind, _pc, target = self.terminals[0]
        if kind == "tail_j" and target is not None:
            return f"term=tail_j:0x{target:08X}"
        return f"term={kind}"

    def escape_evidence(self) -> str:
        assert self.escape_at is not None
        pc, target, kind = self.escape_at
        return f"escape={kind}@0x{pc:08X}->0x{target:08X}"


def signed16(value: int) -> int:
    return value - 0x10000 if value & 0x8000 else value


def jump_target(pc: int, word: int) -> int:
    return ((pc + 4) & 0xF0000000) | ((word & 0x03FFFFFF) << 2)


def branch_target(pc: int, word: int) -> int:
    return pc + 4 + (signed16(word & 0xFFFF) << 2)


def is_branch(word: int) -> bool:
    opcode = word >> 26
    if opcode in (0x04, 0x05, 0x06, 0x07):
        return True
    if opcode == 0x01:                       # REGIMM: bltz/bgez/bltzal/bgezal (+ likely)
        return ((word >> 16) & 0x1F) in (0, 1, 2, 3, 16, 17, 18, 19)
    if opcode == 0x10:                       # COP0: bc0f/bc0t
        return ((word >> 21) & 0x1F) == 8
    return False


def walk(payload: bytes, text_address: int, start: int, hard: frozenset[int],
         hard_sorted: Sequence[int]) -> Walk:
    """Explore every path from `start` and report its reachable extent."""
    limit = text_address + len(payload)
    pending = [start]
    visited: set[int] = set()
    reach = start - 4
    boundary: int | None = None
    indirect_at: int | None = None
    terminals: list[tuple[str, int, int | None]] = []
    escape_at: tuple[int, int, str] | None = None
    overflow: tuple[str, int] | None = None

    def next_hard_after(pc: int) -> int | None:
        index = bisect.bisect_right(hard_sorted, pc)
        return hard_sorted[index] if index < len(hard_sorted) else None

    def note_escape(pc: int, target: int, kind: str) -> None:
        nonlocal escape_at
        following = next_hard_after(pc)
        if escape_at is None and following is not None and target > following:
            escape_at = (pc, target, kind)

    while pending:
        pc = pending.pop()
        while True:
            if pc in visited:
                break
            if pc < text_address or pc + 4 > limit:
                overflow = ("outside", pc)
                break
            if pc - start > MAX_FUNCTION_SIZE:
                overflow = ("runaway", pc)
                break
            if pc != start and pc in hard:
                if boundary is None or pc < boundary:
                    boundary = pc
                break
            visited.add(pc)
            if pc + 4 > reach:
                reach = pc + 4
            word = struct.unpack_from("<I", payload, pc - text_address)[0]
            opcode = word >> 26
            funct = word & 0x3F

            if opcode == 0 and funct == 0x08:                 # jr
                if pc + 8 > reach:
                    reach = pc + 8                            # the delay slot is part of the body
                rs = (word >> 21) & 0x1F
                if rs == 31:
                    terminals.append(("jr_ra", pc, None))
                else:
                    if indirect_at is None:
                        indirect_at = pc
                break
            if opcode == 0 and funct == 0x09:                 # jalr
                pending.append(pc + 4)
                pc += 8
                continue
            if opcode == 0 and funct == 0x0D:                 # break (traps)
                terminals.append(("break", pc, None))
                break
            if opcode == 0 and funct == 0x0C:                 # syscall (returns)
                pc += 4
                continue
            if opcode == 0x02:                                # j
                target = jump_target(pc, word)
                # The delay slot executes and is part of the body, but control
                # then leaves this instruction; it must not continue linearly.
                if pc + 8 > reach:
                    reach = pc + 8
                if target < start or target in hard:
                    terminals.append(("tail_j", pc, target))
                    break
                note_escape(pc, target, "j")
                pending.append(target)
                break
            if opcode == 0x03:                                # jal
                pending.append(pc + 4)
                pc += 8
                continue
            if is_branch(word):
                target = branch_target(pc, word)
                note_escape(pc, target, "branch")
                pending.append(pc + 4)
                pending.append(target)
                break
            pc += 4
    return Walk(reach, boundary, indirect_at, terminals, escape_at, overflow)


# --------------------------------------------------------------------------
# extent derivation
# --------------------------------------------------------------------------


class Extent:
    """A derived extent (or a classified non-function candidate)."""

    __slots__ = ("start", "end", "next", "grade", "evidence")

    def __init__(self, start: int, end: int | None, next_start: int | None,
                 grade: str, evidence: list[str]) -> None:
        self.start = start
        self.end = end
        self.next = next_start
        self.grade = grade
        self.evidence = evidence

    @property
    def size(self) -> int | None:
        return None if self.end is None else self.end - self.start

    @property
    def gap(self) -> int | None:
        if self.end is None or self.next is None:
            return None
        return self.next - self.end


def derive_extents(payload: bytes, text_address: int,
                   grades: dict[int, set[str]]) -> list[Extent]:
    """Derive an extent for every candidate start, in address order."""
    hard = frozenset(a for a, names in grades.items() if names & HARD_GRADES)
    hard_sorted = sorted(hard)
    following: dict[int, int | None] = {}
    for index, address in enumerate(hard_sorted):
        following[address] = (hard_sorted[index + 1] if index + 1 < len(hard_sorted) else None)

    extents: dict[int, Extent] = {}
    for start in hard_sorted:
        next_start = following[start]
        result = walk(payload, text_address, start, hard, hard_sorted)
        evidence: list[str] = []

        # A hard start is another function's entry, so this function's extent
        # cannot extend past it; `clip` is the boundary the result is bounded by.
        clip = next_start if next_start is not None else result.boundary
        escaped = result.escape_at is not None or (clip is not None and result.reach > clip)

        if result.overflow is not None:
            kind, pc = result.overflow
            grade = kind
            end = clip if clip is not None else result.reach
            evidence.append(f"pc=0x{pc:08X}")
        elif escaped:
            grade = "escape"
            end = clip if clip is not None else result.reach
            if result.escape_at is not None:
                evidence.append(result.escape_evidence())
            else:
                evidence.append(f"over=0x{result.reach:08X}")
        elif result.boundary is not None:
            grade = "fallthrough"
            end = result.boundary
            evidence.append(f"hit=0x{result.boundary:08X}")
        elif result.indirect_at is not None:
            grade = "indirect"
            end = next_start if next_start is not None else result.reach
            evidence.append(f"jr_reg=0x{result.indirect_at:08X}")
            evidence.append(f"lower_bound=0x{result.reach:08X}")
        else:
            grade = "exact"
            end = result.reach
            evidence.append(result.terminal_evidence())
        extents[start] = Extent(start, end, next_start, grade, evidence)

    # Soft candidates: containment inside an `exact` extent is evidence that the
    # candidate is not an independent function. Extents of other grades are not
    # trustworthy enough to carry that inference.
    exact_spans = sorted(
        (extent.start, extent.end) for extent in extents.values()
        if extent.grade == "exact" and extent.end is not None
    )

    def contained_by(address: int) -> int | None:
        for span_start, span_end in exact_spans:
            if span_start < address < span_end:
                return span_start
            if span_start > address:
                break
        return None

    for address in sorted(set(grades) - set(hard)):
        names = ",".join(sorted(grades[address]))
        owner = contained_by(address)
        if owner is None:
            extents[address] = Extent(address, None, None, "standalone", [f"grades={names}"])
        else:
            extents[address] = Extent(address, None, None, "contained",
                                      [f"inside=0x{owner:08X}", f"grades={names}"])

    return [extents[address] for address in sorted(extents)]


def format_extents(extents: Sequence[Extent]) -> str:
    lines = list(HEADER_LINES)
    for extent in extents:
        end = "-" if extent.end is None else f"0x{extent.end:08X}"
        size = "-" if extent.size is None else str(extent.size)
        next_start = "-" if extent.next is None else f"0x{extent.next:08X}"
        gap = "-" if extent.gap is None else str(extent.gap)
        evidence = ";".join(extent.evidence)
        lines.append("\t".join((
            f"0x{extent.start:08X}", end, size, next_start, gap, extent.grade, evidence,
        )))
    return "\n".join(lines) + "\n"


# --------------------------------------------------------------------------
# subcommands
# --------------------------------------------------------------------------


def load_payload(exe_path: Path) -> tuple[bytes, int]:
    data = exe_path.read_bytes()
    _entry, text_address, text_size = parse_psx_exe(data)
    payload = data[PAYLOAD_LMA:PAYLOAD_LMA + text_size]
    if len(payload) != text_size:
        raise ToolError("executable is truncated: payload is shorter than the header declares")
    return payload, text_address


def command_scan(args: argparse.Namespace) -> int:
    exe_path = require_file(args.exe, "executable")
    inventory_path = require_file(args.inventory, "function inventory")
    out = resolve_output(args.out, args.force)
    payload, text_address = load_payload(exe_path)
    grades = load_inventory(inventory_path)
    extents = derive_extents(payload, text_address, grades)

    counts = {grade: 0 for grade in GRADES}
    for extent in extents:
        counts[extent.grade] = counts.get(extent.grade, 0) + 1
    out.parent.mkdir(parents=True, exist_ok=True)
    out.write_text(format_extents(extents), encoding="ascii")

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


def command_verify(args: argparse.Namespace) -> int:
    exe_path = require_file(args.exe, "executable")
    inventory_path = require_file(args.inventory, "function inventory")
    extents_path = require_file(args.extents, "function extents")
    regions_path = require_file(args.regions, "region registry")

    payload, text_address = load_payload(exe_path)
    grades = load_inventory(inventory_path)
    derived = {extent.start: extent for extent in derive_extents(payload, text_address, grades)}
    stored = load_inventory_extents(extents_path)
    regions = load_regions(regions_path)

    failures: list[str] = []
    for start, end, source in regions:
        if start not in stored:
            failures.append(f"0x{start:08X} ({source}): absent from the extents table")
            continue
        if stored[start] != end:
            failures.append(
                f"0x{start:08X} ({source}): extents table says 0x{stored[start]:08X}, "
                f"registry says 0x{end:08X}")
            continue
        extent = derived.get(start)
        if extent is None:
            failures.append(f"0x{start:08X} ({source}): not a candidate in the inventory")
            continue
        if extent.end != end:
            failures.append(
                f"0x{start:08X} ({source}): derived 0x{extent.end:08X} (grade {extent.grade}), "
                f"registry says 0x{end:08X}")

    for failure in failures:
        print(f"disagreement: {failure}", file=sys.stderr)
    if failures:
        print(f"regions={len(regions)} disagreements={len(failures)}")
        return 1
    print(f"regions={len(regions)} disagreements=0")
    print("result=AGREE")
    return 0


def load_inventory_extents(path: Path) -> dict[int, int]:
    """Read back a generated extents table: address -> end (function rows only)."""
    ends: dict[int, int] = {}
    for number, raw in enumerate(path.read_text(encoding="utf-8").splitlines(), 1):
        line = raw.split("#", 1)[0].strip()
        if not line:
            continue
        fields = line.split("\t")
        if len(fields) != 7:
            raise ToolError(f"extents line {number}: expected seven fields")
        address = parse_hex(fields[0], f"extents line {number}")
        if fields[1] == "-":
            continue
        ends[address] = parse_hex(fields[1], f"extents line {number}")
    return ends


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="derive extents for every candidate")
    scan_parser.add_argument("--exe", required=True, type=Path)
    scan_parser.add_argument("--inventory", required=True, type=Path,
                             help="graded function-boundary inventory (address<TAB>grades)")
    scan_parser.add_argument("--out", required=True, type=Path)
    scan_parser.add_argument("--force", action="store_true",
                             help="overwrite an existing output file")
    scan_parser.set_defaults(handler=command_scan)

    verify_parser = subparsers.add_parser(
        "verify", help="check derived extents against the registered regions")
    verify_parser.add_argument("--exe", required=True, type=Path)
    verify_parser.add_argument("--inventory", required=True, type=Path)
    verify_parser.add_argument("--extents", required=True, type=Path)
    verify_parser.add_argument("--regions", required=True, type=Path)
    verify_parser.set_defaults(handler=command_verify)

    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())
