#!/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.
           `--scheduler` adds conservative classifications for likely
           branch/delay-slot, load-delay, frame, address, loop, and register
           residuals. `--format json` emits the complete diff and optional
           scheduler findings as one machine-readable JSON document; the default
           text output remains backward-compatible.
  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 json
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


# --------------------------------------------------------------------------
# scheduler diagnostics
# --------------------------------------------------------------------------

# These classifiers are intentionally conservative. They identify a likely
# source-level cause; they do not claim that a candidate is correct or incorrect
# beyond the byte comparison already performed by this tool.

MIPS_BRANCH_OPCODES = {0x02, 0x03, 0x04, 0x05, 0x06, 0x07}
MIPS_CONDITIONAL_BRANCH_OPCODES = {0x04, 0x05, 0x06, 0x07}
MIPS_LOAD_OPCODES = {0x20, 0x21, 0x23, 0x24, 0x25}
SCHEDULER_PRIORITY = (
    "length_mismatch",
    "branch_polarity",
    "branch_target",
    "branch_shape",
    "delay_slot",
    "load_delay",
    "frame_epilogue",
    "address_materialization",
    "loop_schedule",
    "register_allocation",
    "immediate_or_field",
    "opcode_mismatch",
    "nop_difference",
    "unknown",
)


@dataclass(frozen=True)
class SchedulerFinding:
    """One classified difference between two instruction streams."""

    kind: str
    address: int
    original: str
    candidate: str
    detail: str
    original_word: int | None
    candidate_word: int | None

    def as_dict(self) -> dict[str, object]:
        def word(value: int | None) -> str | None:
            return None if value is None else f"0x{value:08X}"

        return {
            "address": f"0x{self.address:08X}",
            "kind": self.kind,
            "original": self.original,
            "candidate": self.candidate,
            "detail": self.detail,
            "original_word": word(self.original_word),
            "candidate_word": word(self.candidate_word),
        }


def _word_at(data: bytes, index: int) -> int | None:
    offset = index * 4
    if offset < 0 or offset + 4 > len(data):
        return None
    return int.from_bytes(data[offset:offset + 4], "little")


def _opcode(word: int | None) -> int | None:
    return None if word is None else word >> 26


def _is_branch(word: int | None) -> bool:
    return _opcode(word) in MIPS_BRANCH_OPCODES


def _is_conditional_branch(word: int | None) -> bool:
    return _opcode(word) in MIPS_CONDITIONAL_BRANCH_OPCODES


def _is_load(word: int | None) -> bool:
    return _opcode(word) in MIPS_LOAD_OPCODES


def _is_lui(word: int | None) -> bool:
    return _opcode(word) == 0x0F


def _is_nop(word: int | None) -> bool:
    return word == 0


def _sign_extend(value: int, bits: int) -> int:
    sign = 1 << (bits - 1)
    return (value & (sign - 1)) - (value & sign)


def _branch_target(word: int, pc: int) -> int | None:
    opcode = word >> 26
    if opcode in (0x02, 0x03):
        return ((pc + 4) & 0xF0000000) | ((word & 0x03FFFFFF) << 2)
    if opcode in MIPS_CONDITIONAL_BRANCH_OPCODES:
        return pc + 4 + (_sign_extend(word & 0xFFFF, 16) << 2)
    return None


def _has_backward_branch(data: bytes, start: int) -> bool:
    for index in range(len(data) // 4):
        word = _word_at(data, index)
        if word is None or not _is_conditional_branch(word):
            continue
        pc = start + index * 4
        target = _branch_target(word, pc)
        if target is not None and target <= pc:
            return True
    return False


def _register_fields(word: int) -> tuple[int, ...]:
    opcode = word >> 26
    rs = (word >> 21) & 0x1F
    rt = (word >> 16) & 0x1F
    if opcode == 0:
        return (rs, rt, (word >> 11) & 0x1F, (word >> 6) & 0x1F, word & 0x3F)
    return (rs, rt)


def _looks_like_frame(text: str) -> bool:
    text = text.lower()
    return bool(re.search(r"\bsp\b", text) and
                re.search(r"\b(addiu|lw|sw|jr)\b", text))


def _classify_word_difference(
    index: int,
    start: int,
    original_word: int,
    candidate_word: int,
    original: bytes,
    candidate: bytes,
    original_text: dict[int, str],
    candidate_text: dict[int, str],
) -> tuple[str, str]:
    """Return (kind, human-readable explanation) for one differing word."""
    address = start + index * 4
    a_text = original_text.get(address, "?")
    b_text = candidate_text.get(address, "?")
    a_opcode = _opcode(original_word)
    b_opcode = _opcode(candidate_word)

    if (_is_conditional_branch(original_word) and
            _is_conditional_branch(candidate_word)):
        if a_opcode != b_opcode:
            return (
                "branch_polarity",
                f"conditional branch opcode changed ({a_text} vs {b_text})",
            )
        if (original_word & 0xFFFF) != (candidate_word & 0xFFFF):
            return ("branch_target", "conditional branch target immediate changed")

    if _is_branch(original_word) != _is_branch(candidate_word):
        return ("branch_shape", "one side has a branch where the other does not")

    previous_a = _word_at(original, index - 1)
    previous_b = _word_at(candidate, index - 1)
    if _is_branch(previous_a) or _is_branch(previous_b):
        return (
            "delay_slot",
            "instruction follows a branch and occupies its delay slot",
        )

    if _is_load(previous_a) or _is_load(previous_b):
        return (
            "load_delay",
            "instruction follows a load and may be a load-delay boundary",
        )

    if _looks_like_frame(a_text) or _looks_like_frame(b_text):
        return ("frame_epilogue", "difference is in stack/return-address handling")

    if _is_lui(original_word) != _is_lui(candidate_word):
        return (
            "address_materialization",
            "one side materializes an address with lui while the other does not",
        )

    if (_has_backward_branch(original, start) or
            _has_backward_branch(candidate, start)):
        return ("loop_schedule", "difference occurs in a range containing a backward branch")

    if a_opcode == b_opcode:
        if _register_fields(original_word) != _register_fields(candidate_word):
            return ("register_allocation", "same opcode uses different registers")
        if (original_word & 0xFFFF) != (candidate_word & 0xFFFF):
            return ("immediate_or_field", "same opcode has a different immediate/field")

    if _is_nop(original_word) != _is_nop(candidate_word):
        return ("nop_difference", "one side is a nop where the other is not")
    if a_opcode != b_opcode:
        return ("opcode_mismatch", "instruction opcode changed")
    return ("unknown", "instruction differs but no scheduler rule matched")


def classify_scheduler(
    original: bytes,
    candidate: bytes,
    start: int,
    original_text: dict[int, str] | None = None,
    candidate_text: dict[int, str] | None = None,
) -> list[SchedulerFinding]:
    """Classify byte differences without changing the match verdict.

    The function is pure and accepts synthetic byte streams, which keeps the
    classifier testable without an extracted executable or a local objdump.
    """
    original_text = original_text or {}
    candidate_text = candidate_text or {}
    count = max(len(original), len(candidate)) // 4
    findings: list[SchedulerFinding] = []
    for index in range(count):
        address = start + index * 4
        a_word = _word_at(original, index)
        b_word = _word_at(candidate, index)
        if a_word is None or b_word is None:
            findings.append(SchedulerFinding(
                "length_mismatch",
                address,
                original_text.get(address, "?"),
                candidate_text.get(address, "?"),
                "one instruction stream ends before the other",
                a_word,
                b_word,
            ))
            continue
        if a_word == b_word:
            continue
        kind, detail = _classify_word_difference(
            index,
            start,
            a_word,
            b_word,
            original,
            candidate,
            original_text,
            candidate_text,
        )
        findings.append(SchedulerFinding(
            kind,
            address,
            original_text.get(address, "?"),
            candidate_text.get(address, "?"),
            detail,
            a_word,
            b_word,
        ))
    return findings


def scheduler_summary(findings: Sequence[SchedulerFinding]) -> dict[str, object]:
    if not findings:
        return {"primary": "none", "findings": []}
    priority = {kind: i for i, kind in enumerate(SCHEDULER_PRIORITY)}
    primary = min((f.kind for f in findings), key=lambda kind: priority.get(kind, 999))
    return {
        "primary": primary,
        "findings": [finding.as_dict() for finding in findings],
    }


def print_scheduler_text(findings: Sequence[SchedulerFinding]) -> None:
    summary = scheduler_summary(findings)
    print(f"scheduler_primary={summary['primary']}")
    print(f"scheduler_findings={len(findings)}")
    for finding in findings:
        print(
            f"scheduler_finding=0x{finding.address:08X} "
            f"kind={finding.kind} original={finding.original!r} "
            f"candidate={finding.candidate!r} detail={finding.detail!r}"
        )


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


def _diff_rows(
    args: argparse.Namespace,
    original: bytes,
    candidate: bytes,
    original_disasm: dict[int, str],
    candidate_disasm: dict[int, str],
    registry: dict[int, str],
) -> list[dict[str, object]]:
    rows: list[dict[str, object]] = []
    for off in range(0, max(len(original), len(candidate)), 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(
            original_disasm.get(addr, "?"), registry, args.gp)
        text_b, notes_b = annotate_text(
            candidate_disasm.get(addr, "?"), registry, args.gp)
        rows.append({
            "address": f"0x{addr:08X}",
            "original": text_a,
            "candidate": text_b,
            "notes": list(dict.fromkeys(notes_a + notes_b)),
        })
    return rows


def _json_diff_payload(
    args: argparse.Namespace,
    original: bytes,
    candidate: bytes,
    diff: Sequence[int],
    rows: Sequence[dict[str, object]],
    scheduler: Sequence[SchedulerFinding],
) -> dict[str, object]:
    payload: dict[str, object] = {
        "range": {
            "start": f"0x{args.start:08X}",
            "end": f"0x{args.end:08X}",
        },
        "orig_len": len(original),
        "cand_len": len(candidate),
        "differing_bytes": len(diff),
        "first_difference": (
            None if not diff else
            f"0x{args.start + (diff[0] // 4) * 4:08X}"
        ),
        "result": "MATCH" if not diff else "DIFF",
        "rows": list(rows),
    }
    if args.scheduler:
        payload["scheduler"] = scheduler_summary(scheduler)
    return payload


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]]
    scheduler = classify_scheduler(original, candidate, args.start, da, db) \
        if args.scheduler else []

    if args.format == "json":
        print(json.dumps(
            _json_diff_payload(args, original, candidate, diff,
                               _diff_rows(args, original, candidate, da, db, registry),
                               scheduler),
            sort_keys=True,
        ))
        return 0 if not diff else 1

    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")
        if args.scheduler:
            print_scheduler_text(scheduler)
        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))
    if args.scheduler:
        print_scheduler_text(scheduler)
    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.add_argument("--scheduler", action="store_true",
                        help="classify likely scheduler/live-range residuals")
    p_diff.add_argument("--format", choices=("text", "json"), default="text",
                        help="output format (default: text; JSON is machine-readable)")
    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())