#!/usr/bin/env python3
"""Correctness gate for the search fast paths (run by .auto/checks.sh).

`SearchMixin` grew two optimisations that are invisible to the scenario suite
because they produce the *same answer* when they work:

* **prefix narrowing** — typing a character onto the term rescans only the
  previous hits, because a line holding "mov" holds "mo";
* **the joined haystack** — the whole body is concatenated once so a term is
  found with a C-level `str.find` walk instead of a python loop over every row.

Both are cache-shaped, so the way they break is *staleness*, not a crash. This
drives the real `ListingView` and `DecompView` and asserts that, for every
prefix of a set of terms, the fast path returns exactly the matches and
highlight ranges the plain per-line loop does — including after the things that
are meant to invalidate them (ending a search, toggling the opcode column,
navigating).

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

import asyncio
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, os.path.join(ROOT, "tests"))
sys.path.insert(0, HERE)

from bench import stage  # noqa: E402
from idatui._sync import wait_for  # noqa: E402
from idatui.app import DecompView, IdaTui, ListingView  # noqa: E402

TERMS = ("mov", "call", "rsp", "Mov", "1a", "push", "e", "lea", "sub_", "0",
         "  ", "if", "v1", "]")


def _reference(view, term: str):
    """(matches, ranges) from the plain per-line loop, with every cache off."""
    cls = view.__class__
    saved = cls._search_haystack
    cls._search_haystack = lambda self, c, s: None
    try:
        view._reset_search_cache(body=True)
        view._term = term
        view._ci = term.islower()
        view._compute_matches()
        return (list(view._matches),
                {k: list(v) for k, v in view._ranges.items()})
    finally:
        cls._search_haystack = saved
        view._reset_search_cache(body=True)


def _fast(view, term: str):
    view._term = term
    view._ci = term.islower()
    view._compute_matches()
    return (list(view._matches), {k: list(v) for k, v in view._ranges.items()})


def _check(view, name: str, fails: list) -> int:
    """Type every prefix of every term; compare the fast path to the loop."""
    checked = 0
    for word in TERMS:
        view.search_begin(1)
        for i in range(1, len(word) + 1):
            term = word[:i]
            got = _fast(view, term)         # may narrow from the previous term
            want = _reference(view, term)
            checked += 1
            if got != want:
                a, b = set(got[0]), set(want[0])
                fails.append(
                    f"{name} {term!r}: fast={len(got[0])} loop={len(want[0])} "
                    f"only-fast={sorted(a - b)[:4]} only-loop={sorted(b - a)[:4]}")
        view.search_cancel()
    return checked


async def main() -> int:
    target = sys.argv[1] if len(sys.argv) > 1 else "targets/echo"
    fails: list[str] = []
    d, path = stage(os.path.join(ROOT, target))
    app = IdaTui(open_path=path, keepalive=False)
    try:
        async with app.run_test(size=(140, 44)) as pilot:
            async def w(pred, t=300.0):
                return await wait_for(pred, pilot.pause, t, 0.02)

            if not await w(lambda: app._func_index is not None
                           and app._func_index.complete):
                fails.append("boot: function index never completed")
                return 1
            await w(lambda: app._loading_screen is None
                    and len(app.screen_stack) == 1, 120)
            funcs = app._func_index.all_loaded()
            lst = app.query_one(ListingView)
            fn = funcs[min(5, len(funcs) - 1)]
            app._open_function(fn.addr, fn.name)
            await w(lambda: lst.total > 0 and lst._cursor_ea() == fn.addr, 120)
            lst.model.load_all()
            lst.total = len(lst.model)

            n = _check(lst, "listing", fails)

            # The opcode-bytes column is searchable text and changes width
            # without changing the row count or the model -- the one thing the
            # haystack's key cannot see.
            app.action_toggle_view.__self__  # noqa: B018 - keep app referenced
            lst.action_toggle_opcodes()
            n += _check(lst, "listing/opcodes", fails)
            lst.action_toggle_opcodes()

            # Navigating inside the same segment must NOT invalidate the body,
            # and must not leave it stale either.
            other = funcs[min(9, len(funcs) - 1)]
            app._open_function(other.addr, other.name)
            await w(lambda: lst.total > 0
                    and lst._cursor_ea() == other.addr, 120)
            lst.model.load_all()
            lst.total = len(lst.model)
            n += _check(lst, "listing/after-nav", fails)

            # The pseudocode view uses a different line source.
            app._active = "listing"
            app._show_active()
            lst.focus()
            await pilot.pause(0.05)
            app.action_toggle_view()
            dv = app.query_one(DecompView)
            if await w(lambda: dv.total > 0, 60):
                n += _check(dv, "decomp", fails)
            app.exit()
        print(f"search fast paths: {n} prefixes checked, {len(fails)} mismatches")
        for f in fails[:10]:
            print("  FAIL", f)
        return 1 if fails else 0
    finally:
        shutil.rmtree(d, ignore_errors=True)


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