#!/usr/bin/env python3
"""Differential gate for the fast decompile_function_safe (run by checks.sh).

`idatui/worker.py` rebinds ida-pro-mcp's `decompile_function_safe` to our own
loop, which skips two of the three SWIG allocations per line and memoises the
per-line `dstr()` by ctree obj_id. That is a pure speed change and the text it
returns is what the pseudocode pane shows, markers and all -- so it has to be
byte-identical, not merely similar.

This runs BOTH implementations against the same cfunc for every function of a
real binary and compares the strings, with `include_addresses` both ways (the
marker path is the whole point, and the no-marker path must not regress either).

    ~/ida-venv/bin/python .auto/check_decomp.py [targets/echo] [max_funcs]
"""
from __future__ import annotations

import os
import shutil
import sys

HERE = os.path.dirname(os.path.abspath(__file__))
ROOT = os.path.dirname(HERE)
sys.path.insert(0, ROOT)
sys.path.insert(0, HERE)

from bench import stage  # noqa: E402


def _add_mcp_path() -> None:
    """ida-pro-mcp is installed for the interpreter the WORKER runs, which is
    not necessarily the one running this check."""
    import glob
    for pat in ("/home/user/.local/lib/python3.*/site-packages",
                os.path.expanduser("~/.local/lib/python3.*/site-packages")):
        for d in glob.glob(pat):
            if os.path.isdir(os.path.join(d, "ida_pro_mcp")) and d not in sys.path:
                sys.path.append(d)


def main() -> int:
    target = sys.argv[1] if len(sys.argv) > 1 else "targets/echo"
    limit = int(sys.argv[2]) if len(sys.argv) > 2 else 400
    d, path = stage(os.path.join(ROOT, target))
    os.environ["IDA_MCP_TOOL_TIMEOUT_SEC"] = "0"
    _add_mcp_path()
    try:
        import idapro
        idapro.open_database(path, run_auto_analysis=True)
        try:
            import ida_funcs
            import ida_hexrays
            import idautils
            from ida_pro_mcp.ida_mcp import utils

            original = utils.decompile_function_safe
            sys.path.insert(0, os.path.join(ROOT, "server"))
            import patch_server
            # exec only our function out of the decoded BODY: the rest of it
            # needs api_types' namespace (@tool, @idasync, ...).
            body = patch_server.BODY
            i = body.find("def _idatui_decompile_function_safe(")
            j = body.find("\n_idatui_strings_cache", i)
            assert i > 0 and j > i, "could not slice the function out of BODY"
            g = {}
            exec(body[i:j], g)
            fast = g["_idatui_decompile_function_safe"]

            ida_hexrays.init_hexrays_plugin()
            fails: list[str] = []
            n = ok = 0
            for ea in idautils.Functions():
                f = ida_funcs.get_func(ea)
                if not f:
                    continue
                n += 1
                if n > limit:
                    break
                for markers in (True, False):
                    a, ea_err = original(f.start_ea, include_addresses=markers)
                    b, eb_err = fast(f.start_ea, include_addresses=markers)
                    if a != b or (ea_err is None) != (eb_err is None):
                        fails.append(f"{f.start_ea:#x} (markers={markers})")
                        if len(fails) <= 3:
                            la = (a or "").splitlines()
                            lb = (b or "").splitlines()
                            i = next((k for k, (x, y) in enumerate(zip(la, lb))
                                      if x != y), None)
                            if i is None:
                                print(f"  FAIL {f.start_ea:#x}: {len(la)} lines "
                                      f"vs {len(lb)}, errs {ea_err!r}/{eb_err!r}")
                            else:
                                print(f"  FAIL {f.start_ea:#x} line {i}:")
                                print(f"      original: {la[i][:100]!r}")
                                print(f"      fast    : {lb[i][:100]!r}")
                        break
                else:
                    ok += 1
            print(f"decompile text: {ok}/{min(n, limit)} functions of {target} "
                  f"byte-identical to ida-pro-mcp's own loop, "
                  f"{len(fails)} problems")
            return 1 if fails else 0
        finally:
            idapro.close_database(save=False)
    finally:
        shutil.rmtree(d, ignore_errors=True)


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