#!/usr/bin/env python3
"""Duplicate-body census for the USA executable.

Matching conventions require checking for duplicates *before* registering a
match, because a body shared by several addresses is matched **once** and
registered once per address (`config/regions.tsv`, N rows to one source). Phase 6
did that check by hand and found one shared body; there was no way to know how
many others exist.

This tool hashes the body of every derived extent in
`config/function_extents.tsv` and groups exact duplicates.

## What is grouped

Every extent row -- any grade -- is hashed over exactly the bytes
`[address, end)` that `tools/sf3_extents` derived. Each member keeps its grade so
a consumer can prefer groups whose members are all `exact`. `--exact-only`
restricts the census to `exact` extents, and `--min-size` drops bodies shorter
than N bytes (very short bodies collide for boring reasons, e.g. `jr ra; nop`).

## What is written

A tracked, sorted TSV of **multi-address groups only**:

    group<TAB>size<TAB>count<TAB>addresses<TAB>grades<TAB>flags

  * `group` is a stable label (`g0001`, ...) assigned in `(size, first address)`
    order, so the file is deterministic;
  * `addresses` is a comma-separated list in ascending order;
  * `grades` is the comma-separated grade of each address, in the same order;
  * `flags` is `zero` when the shared body is nothing but zero bytes, otherwise
    `-`. An all-zero "body" is not a function: it is a zero-filled region the
    walk ran through, which happens when a data word decodes as a `jal` whose
    target lands in the zero band. Those groups are recorded, not hidden, so the
    finding stays visible and the worklist can exclude them.

The content hash is computed and **never written**: grouping is the finding, and
addresses, sizes and grades are the only ROM-derived quantities this project
tracks (the same rule `config/function_inventory.tsv` and
`config/function_extents.tsv` follow). No instruction bytes are read into the
output, and the executable is read only from the caller-supplied path.

The census is a measurement, not a match claim: two identical bodies are
evidence of shared code, and nothing more. A shared *tail* is not a duplicate
function, so a group member still has to be verified as a real function start by
the ordinary workflow.

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

from __future__ import annotations

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


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

EXTENT_GRADES = frozenset({"exact", "fallthrough", "indirect", "escape"})

HEADER_LINES = (
    "# Syphon Filter 3 (USA) duplicate-body census.",
    "# Columns: group<TAB>size<TAB>count<TAB>addresses<TAB>grades<TAB>flags.",
    "# Multi-address groups only; addresses, sizes and grades only; no bytes.",
    "# A group is evidence of a shared body, not a claim that every member is",
    "# a real function start: verify a member before registering it.",
    "# flags=zero means the body is all zero bytes -- a zero-filled region the",
    "# walk ran through, not a function. Exclude those from matching.",
    "# Regenerate: ./tools/sf3_dupes census --exe '<exe>' \\",
    "#   --extents config/function_extents.tsv --out config/duplicate_bodies.tsv --force",
)


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 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_extents(path: Path) -> list[tuple[int, int, str]]:
    """Read a generated extents table; keep rows that carry an extent."""
    rows: 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) != 7:
            raise ToolError(f"extents line {number}: expected seven fields")
        address = parse_hex(fields[0], f"extents line {number}")
        grade = fields[5]
        if fields[1] == "-":
            continue
        end = parse_hex(fields[1], f"extents line {number}")
        if end <= address:
            raise ToolError(f"extents line {number}: end is not after address")
        if grade not in EXTENT_GRADES:
            raise ToolError(f"extents line {number}: {grade!r} carries no extent")
        rows.append((address, end, grade))
    if not rows:
        raise ToolError("extents table contains no rows with an extent")
    return rows


class Group:
    """One set of addresses whose bodies are byte-identical."""

    __slots__ = ("size", "members", "all_zero")

    def __init__(self, size: int, all_zero: bool) -> None:
        self.size = size
        self.members: list[tuple[int, str]] = []
        self.all_zero = all_zero


def census(payload: bytes, text_address: int, rows: Sequence[tuple[int, int, str]],
           min_size: int, exact_only: bool) -> tuple[list[Group], int]:
    """Group rows by (size, body bytes). Returns the groups and the singleton count."""
    buckets: dict[tuple[int, bytes], Group] = {}
    for address, end, grade in rows:
        if exact_only and grade != "exact":
            continue
        size = end - address
        if size < min_size:
            continue
        start = address - text_address
        body = payload[start:start + size]
        if len(body) != size:
            raise ToolError(f"extent 0x{address:08X}..0x{end:08X} is outside the payload")
        key = (size, hashlib.sha1(body).digest())
        group = buckets.get(key)
        if group is None:
            group = buckets[key] = Group(size, not any(body))
        group.members.append((address, grade))

    groups = [group for group in buckets.values() if len(group.members) > 1]
    singletons = sum(1 for group in buckets.values() if len(group.members) == 1)
    for group in groups:
        group.members.sort()
    groups.sort(key=lambda group: (group.size, group.members[0][0]))
    return groups, singletons


def format_census(groups: Sequence[Group]) -> str:
    lines = list(HEADER_LINES)
    for index, group in enumerate(groups, 1):
        addresses = ",".join(f"0x{address:08X}" for address, _grade in group.members)
        grades = ",".join(grade for _address, grade in group.members)
        flags = "zero" if group.all_zero else "-"
        lines.append("\t".join((
            f"g{index:04d}", str(group.size), str(len(group.members)), addresses, grades, flags,
        )))
    return "\n".join(lines) + "\n"


def command_census(args: argparse.Namespace) -> int:
    exe_path = require_file(args.exe, "executable")
    extents_path = require_file(args.extents, "function extents")
    out = resolve_output(args.out, args.force)
    if args.min_size < 0:
        raise ToolError("--min-size cannot be negative")

    data = exe_path.read_bytes()
    _entry, text_address, text_size = parse_psx_exe(data)
    payload = data[PAYLOAD_LMA:PAYLOAD_LMA + text_size]
    rows = load_extents(extents_path)
    groups, singletons = census(payload, text_address, rows, args.min_size, args.exact_only)

    out.parent.mkdir(parents=True, exist_ok=True)
    out.write_text(format_census(groups), encoding="ascii")

    addresses = sum(len(group.members) for group in groups)
    duplicated = sum(group.size * (len(group.members) - 1) for group in groups)
    zero_groups = sum(1 for group in groups if group.all_zero)
    print(f"extents={len(rows)}")
    print(f"groups={len(groups)}")
    print(f"groups_zero_body={zero_groups}")
    print(f"groups_with_code={len(groups) - zero_groups}")
    print(f"addresses_in_groups={addresses}")
    print(f"singletons={singletons}")
    print(f"duplicated_bytes={duplicated}")
    if groups:
        print(f"largest_group={max(len(group.members) for group in groups)}")
    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)

    census_parser = subparsers.add_parser("census", help="group byte-identical bodies")
    census_parser.add_argument("--exe", required=True, type=Path)
    census_parser.add_argument("--extents", required=True, type=Path,
                               help="derived function extents (tools/sf3_extents output)")
    census_parser.add_argument("--out", required=True, type=Path)
    census_parser.add_argument("--force", action="store_true",
                               help="overwrite an existing output file")
    census_parser.add_argument("--min-size", type=int, default=1, dest="min_size",
                               help="ignore bodies shorter than N bytes (default 1)")
    census_parser.add_argument("--exact-only", action="store_true",
                               help="census only extents graded exact")
    census_parser.set_defaults(handler=command_census)

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