#!/usr/bin/env python3
"""Differential check: the current `_idatui_spans` vs the one at a git ref.

The span walker turns IDA's colour-tagged disassembly line into (spans, ops).
It is on the hot path of every listing row, so it is worth optimising -- but its
output drives highlighting, operand marking and the cursor's column arithmetic,
so "faster" is only acceptable if it is byte-identical.

This pulls both implementations out of `server/patch_server.py` (the current
working tree, and whatever `--ref` names), runs them over every tagged line of
a real binary, and reports the first disagreement.

    /usr/bin/python3 .auto/diff_spans.py [--ref HEAD] [--target targets/bash]
                                         [--limit 60000]
"""
from __future__ import annotations

import argparse
import os
import shutil
import subprocess
import sys
import tempfile

ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))


def load_impl(path: str, name: str):
    """Exec just the span-walker section of a patch_server.py's injected BODY.

    BODY is a normal (non-raw) triple-quoted string, so the escapes in it are
    resolved once by importing the module -- slicing the file text instead would
    compile the *undecoded* source and silently test a different program (an
    escaped backslash became a literal one, and every 's' in a comment turned
    into a space).
    """
    import importlib.util
    spec = importlib.util.spec_from_file_location(f"_ps_{name}", path)
    mod = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(mod)          # IDA-free at import time
    body = mod.BODY
    a = body.index("def _idatui_head_row")
    b = body.index("def _idatui_struct_member_rows")
    g = {"__name__": name}
    exec(compile(body[a:b], name, "exec"), g)  # noqa: S102
    return g


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--ref", default="HEAD")
    ap.add_argument("--target", default="targets/bash")
    ap.add_argument("--limit", type=int, default=60000)
    a = ap.parse_args()

    d = tempfile.mkdtemp(prefix="diffspans-")
    old_path = os.path.join(d, "patch_server_old.py")
    with open(old_path, "w") as fh:
        fh.write(subprocess.run(
            ["git", "-C", ROOT, "show", f"{a.ref}:server/patch_server.py"],
            capture_output=True, text=True, check=True).stdout)
    gnew = load_impl(os.path.join(ROOT, "server", "patch_server.py"), "new")
    gold = load_impl(old_path, "old")
    new, old = gnew["_idatui_spans"], gold["_idatui_spans"]
    new_row, old_row = gnew["_idatui_head_row"], gold["_idatui_head_row"]

    binary = os.path.join(ROOT, a.target)
    tgt = os.path.join(d, os.path.basename(binary))
    shutil.copy2(binary, tgt)
    shutil.copy2(binary + ".i64", tgt + ".i64")
    try:
        import idapro
        if idapro.open_database(tgt, run_auto_analysis=False) != 0:
            raise SystemExit("could not open database")
        import ida_bytes
        import ida_lines
        import ida_segment
        import idaapi

        checked = bad = 0
        for si in range(ida_segment.get_segm_qty()):
            seg = ida_segment.getnseg(si)
            if seg is None:
                continue
            ea = seg.start_ea
            while ea < seg.end_ea and checked < a.limit:
                line = ida_lines.generate_disasm_line(ea, 0)
                if line:
                    checked += 1
                    ra, rb = old(line), new(line)
                    if ra != rb:
                        bad += 1
                        if bad <= 3:
                            print(f"SPAN MISMATCH @ {ea:#x}\n  line={line!r}\n"
                                  f"  old={ra!r}\n  new={rb!r}")
                    # The whole row, not just the spans: `text`, the spans/text
                    # agreement guard and the name all moved around too.
                    ra, rb = old_row(ea), new_row(ea)
                    if ra != rb:
                        bad += 1
                        if bad <= 3:
                            print(f"ROW MISMATCH @ {ea:#x}\n"
                                  f"  old={ra!r}\n  new={rb!r}")
                nxt = ida_bytes.get_item_end(ea)
                ea = nxt if nxt > ea else ea + 1
            if checked >= a.limit:
                break
        print(f"checked {checked} lines, {bad} mismatches")
        idapro.close_database(False)
        return 1 if bad else 0
    finally:
        shutil.rmtree(d, ignore_errors=True)


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