#!/usr/bin/env python3
"""Correctness gate for the listing after an item edit (run by .auto/checks.sh).

An item edit (`c`/`d`/`u`/`p`) changes structure, but only locally: every head
in front of it keeps its address and its row number. So `Program.bump_items(ea)`
keeps the walk up to there instead of discarding the model — worth 257x on a big
binary (4.9s to make one byte into data, against 19ms).

Keeping *anything* across a structural edit is the risky half of that, and it
fails silently: the pane shows rows that are no longer what the database says.
So this drives real edits and compares the kept model against one built from
scratch, row for row — narrow reads (what painting does) and wide ones (what
building the search body does).

The staged database is a throwaway copy and is never saved, so the edits here do
not need undoing and can be as destructive as they like.

    ~/ida-venv/bin/python .auto/check_edit.py [targets/echo]
"""
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
from idatui.domain import Program  # noqa: E402
from idatui.worker_client import WorkerClient  # noqa: E402


def snapshot(model, total: int, wide: bool):
    if wide:
        rows = []
        for base in range(0, total, 4096):
            rows.extend(model.window(base, min(4096, total - base)))
    else:
        rows = [model.get(i) for i in range(total)]
    return [(h.ea, h.kind, h.size, h.text, h.name) if h else None for h in rows]


def main() -> int:
    target = sys.argv[1] if len(sys.argv) > 1 else "targets/echo"
    d, path = stage(os.path.join(ROOT, target))
    fails: list[str] = []
    client = WorkerClient(path)
    try:
        prog = Program(client)
        idx = prog.functions()
        idx.load_all()
        funcs = sorted(idx.all_loaded(), key=lambda f: f.addr)
        if len(funcs) < 8:
            print(f"{target}: too few functions to check")
            return 1
        seg = funcs[len(funcs) // 2].addr
        model = prog.listing(seg)
        model.load_all()
        total = len(model)

        # Edits at a spread of positions: near the end (where truncation keeps
        # nearly everything), the middle, and early on (where it must give up).
        spots = [funcs[int(len(funcs) * f)].addr for f in (0.9, 0.5, 0.05)]
        for k, ea in enumerate(spots):
            # Undefine, at three sizes. This is the edit that can coalesce
            # BACKWARDS into the undefined run in front of it, which is the
            # reason truncate_from drops two pages rather than one.
            size = (1, 4, 16)[k % 3]
            kind = f"undefine {size}B"
            try:
                prog.undefine(ea, size)
            except Exception as e:  # noqa: BLE001
                fails.append(f"{kind} at {ea:#x} failed: {e}")
                continue
            prog.bump_items(ea)

            kept = prog.listing(seg)
            kept_total = None
            got_wide = None
            if kept is not None:
                kept.load_all()
                kept_total = len(kept)
                got_wide = snapshot(kept, kept_total, wide=True)
                got_narrow = snapshot(kept, kept_total, wide=False)

            # ... against a model that knows nothing about what came before.
            prog._listings.clear()
            fresh = prog.listing(seg)
            fresh.load_all()
            want = snapshot(fresh, len(fresh), wide=False)

            if kept_total != len(fresh):
                fails.append(f"{kind} at {ea:#x}: kept model has {kept_total} "
                             f"rows, a rebuild has {len(fresh)}")
            elif got_narrow != want or got_wide != want:
                which = "narrow" if got_narrow != want else "wide"
                bad = next((i for i, (a, b) in
                            enumerate(zip(got_narrow if which == "narrow"
                                          else got_wide, want)) if a != b), None)
                fails.append(
                    f"{kind} at {ea:#x}: {which} read differs from a rebuild at "
                    f"row {bad}: {(got_narrow if which == 'narrow' else got_wide)[bad]}"
                    f" vs {want[bad]}")

            model = prog.listing(seg)
            model.load_all()
            total = len(model)

        print(f"item edits: {len(spots)} edits checked against a rebuild "
              f"({total} rows), {len(fails)} problems")
        for f in fails:
            print("  FAIL", f)
        return 1 if fails else 0
    finally:
        try:
            client.close()
        except Exception:  # noqa: BLE001
            pass
        shutil.rmtree(d, ignore_errors=True)


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