#!/usr/bin/env python3
"""Instruction-level diff and address resolver for one matching attempt.

Phase 8's workers diagnosed near-misses with a scratch `cmp.py` under the
ignored `.run/` tree. Phase 9 promotes it into the tracked toolchain, hardened
as its author specified:

  * the PS-X text address (LOAD) is read from the executable's header, never
    hardcoded, and the payload file offset is derived from it;
  * the objdump path is derived from the repository (`tools/.../binutils`);
  * every temp file is written under a caller-supplied directory, because
    three Phase 9 sessions share one worktree and a fixed `/tmp` name would
    collide;
  * the address resolver prints ready-to-paste symbol names
    (`D_80122308`, `func_80020000`) rather than bare hex, so a worker copies
    the name the harness will resolve.

It has two subcommands:

  diff     side-by-side instruction diff of the original payload range and one
           candidate `.bin`, the same bytes `sf3_match range` links into
           `<work>/single.bin`. Prints the differing instruction count, the
           first differing address, and a per-instruction table with each
           resolved address annotated with a ready-to-paste symbol name.
  resolve  read an objdump-style disassembly listing and print, per line, the
           resolved absolute address of every `off(gp)` operand and every
           `lui r,HI` / `off(r)` pair, again as ready-to-paste names.

Neither subcommand matches anything: `sf3_match range` is the per-function
comparator and `sf3_match gate` the whole-binary authority. This tool exists so
a worker can see *which* instructions differ and *what address each operand
means*, which is what turns a `DIFF` into a decision.

The tool never embeds game bytes in its source. Tests use self-authored
fixtures only and never read the extracted executable.

Exit codes: 0 success, 1 mismatch (diff found), 2 usage or environment error.
"""

from __future__ import annotations

import argparse
import re
from dataclasses import dataclass
from pathlib import Path
import struct
import subprocess
import sys
from typing import Sequence


REPO_ROOT = Path(__file__).resolve().parent.parent
DEFAULT_OBJDUMP = (
    REPO_ROOT
    / "tools/mipsel-none-elf-binutils/prefix/usr/bin/mipsel-none-elf-objdump"
)
HEADER_SIZE = 0x800
PAYLOAD_LMA = 0x800
EXE_MAGIC = b"PS-X EXE"
# The Phase 6 small-data census fixed gp at this address (cookbook finding 10).
DEFAULT_GP = 0x80121938


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


# --------------------------------------------------------------------------
# PS-X EXE parsing (aligned with sf3_match's parse_psx_exe)
# --------------------------------------------------------------------------


@dataclass(frozen=True)
class PsxExe:
    entry: int
    text_address: int
    text_size: int

    @property
    def payload_end(self) -> int:
        return self.text_address + self.text_size

    def file_offset(self, address: int) -> int:
        """File offset of a payload address (the payload starts at 0x800)."""
        return PAYLOAD_LMA + (address - self.text_address)


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


# --------------------------------------------------------------------------
# objdump
# --------------------------------------------------------------------------


def disassemble(
    data: bytes, vaddr: int, objdump: Path, work: Path, stem: str
) -> dict[int, str]:
    """Disassemble raw bytes as MIPS at `vaddr` and return addr -> text rows.

    All temp files go under `work` (caller-supplied, so sessions never
    collide); `stem` keeps several dumps from the same workdir apart.
    """
    binary = work / f"{stem}.bin"
    binary.write_bytes(data)
    out = subprocess.run(
        [
            str(objdump), "-D", "-b", "binary", "-m", "mips:3000",
            f"--adjust-vma=%#x" % vaddr, str(binary),
        ],
        capture_output=True, text=True, check=True,
    ).stdout
    rows: dict[int, str] = {}
    for line in out.splitlines():
        parts = line.split("\t")
        if len(parts) >= 3 and parts[0].endswith(":"):
            try:
                addr = int(parts[0].rstrip(":"), 16)
            except ValueError:
                continue
            rows[addr] = parts[2].split("#")[0].strip()
    return rows


# --------------------------------------------------------------------------
# address resolution
# --------------------------------------------------------------------------

GP_OPERAND = re.compile(r"(-?\w+)\(gp\)")
LUI = re.compile(r"lui\s+(\w+),(-?\w+)")
MEM = re.compile(r"\w+\s+\w+,(-?\w+)\((\w+)\)")


def symbol_name(address: int, registry: dict[int, str]) -> str:
    """Ready-to-paste name for an address: registry name, else an
    address-shaped placeholder. The harness resolves func_/D_/g_/lbl_ + 8 hex
    digits implicitly, so the placeholder is directly usable in a source file.
    """
    if address in registry:
        return registry[address]
    return f"D_{address:08X}"


def annotate_text(
    text: str, registry: dict[int, str], gp: int
) -> tuple[str, list[str]]:
    """Return (text, notes) with every resolvable operand annotated.

    Notes use ready-to-paste symbol names; for section-ambiguous addresses we
    print the `D_` placeholder and the worker picks func_/g_/lbl_ if the bytes
    say so. A registry row always wins over the placeholder.
    """
    notes: list[str] = []
    mg = GP_OPERAND.search(text)
    if mg:
        try:
            off = int(mg.group(1))
        except ValueError:
            off = 0
        notes.append(f"-> {symbol_name(gp + off, registry)} (gp{off:+d})")
    ml = LUI.match(text)
    if ml:
        try:
            hi = int(ml.group(2), 16) << 16
        except ValueError:
            hi = 0
        mo = MEM.search(text)
        if mo and mo.group(2) == ml.group(1):
            try:
                lo = int(mo.group(1))
            except ValueError:
                lo = 0
            if lo >= 0x8000:
                lo -= 0x10000
            notes.append(f"-> {symbol_name(hi + lo, registry)} (abs)")
    return text, notes


def load_registry(path: Path | None) -> dict[int, str]:
    registry: dict[int, str] = {}
    if path is None:
        return registry
    for line_number, line in enumerate(path.read_text(encoding="utf-8").splitlines(), 1):
        if not line.strip() or line.lstrip().startswith("#"):
            continue
        fields = line.split("\t")
        if len(fields) < 2:
            raise ToolError(f"symbols line {line_number}: expected NAME<TAB>address")
        try:
            address = int(fields[1], 16)
        except ValueError as exc:
            raise ToolError(
                f"symbols line {line_number}: not a hex address: {fields[1]!r}"
            ) from exc
        registry.setdefault(address, fields[0])
    return registry


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


def command_diff(args: argparse.Namespace) -> int:
    data = args.exe.read_bytes()
    exe = parse_psx_exe(data)
    if args.start < exe.text_address or args.end > exe.payload_end:
        raise ToolError("requested range leaves the payload")

    if not args.work.is_dir():
        raise ToolError(f"work directory does not exist: {args.work}")
    if not args.candidate.is_file():
        raise ToolError(f"candidate file does not exist: {args.candidate}")
    registry = load_registry(args.symbols)

    original = data[exe.file_offset(args.start): exe.file_offset(args.end)]
    candidate = args.candidate.read_bytes()

    da = disassemble(original, args.start, args.objdump, args.work, "orig")
    db = disassemble(candidate, args.start, args.objdump, args.work, "cand")

    diff = [i for i in range(len(original)) if i >= len(candidate)
            or original[i] != candidate[i]]
    print(f"range=0x{args.start:08X}..0x{args.end:08X}")
    print(f"orig_len={len(original)} cand_len={len(candidate)} "
          f"differing_bytes={len(diff)}")
    if not diff:
        print("result=MATCH")
        return 0
    first = args.start + (diff[0] // 4) * 4
    print(f"first_difference=0x{first:08X}")
    print("result=DIFF")
    print("%-10s %-22s | %-22s  notes" % ("addr", "ORIGINAL", "CANDIDATE"))
    for off in range(0, len(original), 4):
        ea = original[off:off + 4]
        eb = candidate[off:off + 4] if off < len(candidate) else b""
        if ea == eb:
            continue
        addr = args.start + off
        text_a, notes_a = annotate_text(
            da.get(addr, "?"), registry, args.gp)
        text_b, notes_b = annotate_text(
            db.get(addr, "?"), registry, args.gp)
        notes = " ".join(dict.fromkeys(notes_a + notes_b))
        mark = "  <== first" if addr == first else ""
        print("%-10s %-22s | %-22s  %s%s"
              % (hex(addr), text_a, text_b, notes, mark))
    return 1


def command_resolve(args: argparse.Namespace) -> int:
    registry = load_registry(args.symbols)
    line_re = re.compile(r"^\s*([0-9a-f]{8}):")
    prev_hi: dict[str, int] = {}
    for line in args.disasm.read_text(encoding="utf-8").splitlines():
        m = line_re.match(line)
        addr = int(m.group(1), 16) if m else None
        text = line.split("\t")[-1].split("<")[0].strip() if "\t" in line else line.strip()
        note = ""
        mg = GP_OPERAND.search(text)
        if mg:
            try:
                off = int(mg.group(1))
            except ValueError:
                off = 0
            note = f"  -> {symbol_name(args.gp + off, registry)} (gp{off:+d})"
        ml = LUI.match(text)
        if ml:
            try:
                prev_hi[ml.group(1)] = int(ml.group(2), 16) << 16
            except ValueError:
                pass
            mo = MEM.search(text)
            if mo and mo.group(2) == ml.group(1):
                try:
                    lo = int(mo.group(1))
                except ValueError:
                    lo = 0
                if lo >= 0x8000:
                    lo -= 0x10000
                note = f"  -> {symbol_name(prev_hi[ml.group(1)] + lo, registry)} (abs)"
        mo = MEM.search(text)
        if not note and mo and mo.group(2) in prev_hi:
            try:
                lo = int(mo.group(1))
            except ValueError:
                lo = 0
            if lo >= 0x8000:
                lo -= 0x10000
            note = f"  -> {symbol_name(prev_hi[mo.group(2)] + lo, registry)} (abs)"
        if note:
            prefix = f"{addr:08x}: " if addr is not None else ""
            print(f"{prefix}{text}{note}")
    return 0


# --------------------------------------------------------------------------
# CLI
# --------------------------------------------------------------------------


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        prog="sf3_diff",
        description="Instruction-level diff and address resolver for one "
                    "matching attempt. Phase 8's worker scratch helper, "
                    "promoted and hardened in Phase 9.",
    )
    sub = parser.add_subparsers(dest="command", required=True)

    p_diff = sub.add_parser("diff", help="diff one range against a candidate")
    p_diff.add_argument("--exe", required=True, type=Path,
                        help="validated PS-X EXE (the ORIGINAL)")
    p_diff.add_argument("--start", required=True, help="range start, hex")
    p_diff.add_argument("--end", required=True, help="range end, hex")
    p_diff.add_argument("--candidate", required=True, type=Path,
                        help="candidate .bin (sf3_match range writes "
                             "<work>/single.bin)")
    p_diff.add_argument("--work", required=True, type=Path,
                        help="caller-supplied directory for temp files "
                             "(sessions share the worktree; never /tmp)")
    p_diff.add_argument("--objdump", type=Path, default=DEFAULT_OBJDUMP,
                        help="objdump binary (default: repository binutils)")
    p_diff.add_argument("--symbols", type=Path, default=None,
                        help="tracked symbol registry (NAME<TAB>address)")
    p_diff.add_argument("--gp", type=lambda s: int(s, 0), default=DEFAULT_GP,
                        help=f"gp base for gp-relative resolution "
                             f"(default: {DEFAULT_GP:#x})")
    p_diff.set_defaults(handler=command_diff)

    p_res = sub.add_parser("resolve", help="resolve addresses in a disassembly")
    p_res.add_argument("--disasm", required=True, type=Path,
                       help="objdump-style disassembly text")
    p_res.add_argument("--symbols", type=Path, default=None,
                       help="tracked symbol registry (NAME<TAB>address)")
    p_res.add_argument("--gp", type=lambda s: int(s, 0), default=DEFAULT_GP,
                       help=f"gp base (default: {DEFAULT_GP:#x})")
    p_res.set_defaults(handler=command_resolve)
    return parser


def main(argv: Sequence[str] | None = None) -> int:
    parser = build_parser()
    args = parser.parse_args(argv)
    try:
        if getattr(args, "start", None) is not None:
            args.start = int(args.start, 0)
        if getattr(args, "end", None) is not None:
            args.end = int(args.end, 0)
        return args.handler(args)
    except ToolError as exc:
        print(f"sf3_diff: {exc}", file=sys.stderr)
        return 2


if __name__ == "__main__":
    sys.exit(main())