#!/usr/bin/env python3
"""ida-tui performance benchmark.

Boots the real Textual app over a real idalib worker (a staged copy of a target,
seeded from a pristine .i64 so auto-analysis is not paid per run) and times a
scripted, deterministic set of the things a user actually waits on:

    boot           worker spawn + db open + function index complete + usable
    listing        sweeping the unified listing (cold pages, then warm)
    render         render_line throughput on a warm listing
    decomp         F5 -> decompile -> highlight -> paint, per function
    graph          flowchart -> Sugiyama layout -> paint, per function
    search         incremental search over the whole segment
    index          function index load_all (3 filters) + fuzzy palette filter
    hex            hex view scroll sweep
    pure_graph     offline layout over a frozen CFG corpus (no IDA at all)

Two targets are benched so nothing tunes for one scale: a small binary (echo,
~128 funcs) and a big one (bash, ~2100 funcs).

Every phase is deterministic: the function set is "largest by size, ties by
address, skipping any the graph view refuses" -- never "whatever is biggest
right now", which drifts with analysis.

Prints `METRIC name=value` lines; the primary metric is `total_ms`. Exits
non-zero if any phase failed to do its work (a phase that silently no-ops would
otherwise look like a huge speedup).

    ~/ida-venv/bin/python .auto/bench.py
"""
from __future__ import annotations

import argparse
import asyncio
import json
import os
import shutil
import statistics
import sys
import tempfile
import time

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"))

from idatui.app import (  # noqa: E402
    DecompView, GraphView, HexView, IdaTui, ListingView,
)
from idatui._sync import wait_for  # noqa: E402

SIZE = (140, 44)
TIMES: dict[str, list[float]] = {}
NOTES: dict[str, object] = {}
FAILS: list[str] = []
PREFIX = ""


class _T:
    def __init__(self, key):
        self.key = PREFIX + key

    def __enter__(self):
        self.t = time.perf_counter()
        return self

    def __exit__(self, *a):
        TIMES.setdefault(self.key, []).append((time.perf_counter() - self.t) * 1000)
        return False


def note(k, v):
    NOTES[PREFIX + k] = v


def fail(msg):
    FAILS.append(PREFIX + msg)


#: Poll interval for every "has it landed yet" wait. The phases time the wait,
#: so the interval is measurement overhead: at the 10ms this used to use, a
#: graph open that really took 30ms was charged up to 40, and the graph phase
#: (24 timed waits) carried ~12% of pure quantisation.
_STEP = 0.002


async def _wait(pilot, pred, t=60.0, step=_STEP):
    return await wait_for(pred, pilot.pause, t, step)


def _paint(view) -> int:
    """Render every visible row; returns total cells painted."""
    n = 0
    for y in range(view.size.height):
        n += len(view.render_line(y).text)
    return n


# --------------------------------------------------------------------------- #
# phases
# --------------------------------------------------------------------------- #
async def phase_nav(app, pilot, funcs):
    """Jump to each function's entry in the unified listing, cold.

    This is `g <addr>`, follow, and every xref landing: the listing model has to
    have walked heads as far as the target before the cursor can sit on it. The
    order deliberately alternates far and near (last, first, second-last, ...)
    so it measures random access, not one tidy forward sweep.
    """
    lst = app.query_one(ListingView)
    order = []
    lo, hi = 0, len(funcs) - 1
    while lo <= hi:
        order.append(funcs[hi])
        if lo != hi:
            order.append(funcs[lo])
        lo, hi = lo + 1, hi - 1
    worst = 0.0
    with _T("nav_ms"):
        for fn in order:
            t0 = time.perf_counter()
            app._open_function(fn.addr, fn.name)
            if not await _wait(pilot, lambda fn=fn: app._cur is not None
                               and app._cur.ea == fn.addr and lst.total > 0
                               and lst._cursor_ea() == fn.addr, 120):
                fail(f"nav:{fn.name}")
            worst = max(worst, (time.perf_counter() - t0) * 1000)
    note("nav_worst_ms", round(worst, 1))
    note("nav_rows", lst.total)


async def phase_listing(app, pilot, funcs, pages, first):
    """Sweep the unified listing, painting each viewport. The first pass pulls
    cold pages over the worker; the second is all-cached, i.e. pure python."""
    lst = app.query_one(ListingView)
    fn = funcs[0]
    app._open_function(fn.addr, fn.name)
    if not await _wait(pilot, lambda: app._cur is not None and app._cur.ea == fn.addr
                       and lst.total > 0, 60):
        fail("listing:open")
        return
    app._active = "listing"
    app._show_active()
    lst.focus()
    await pilot.pause(0.02)
    height = max(lst.size.height, 1)

    async def sweep():
        cells = 0
        for i in range(pages):
            lst.model.ensure(i * height + height + 1)
            lst.cursor = min(max(lst.total - 1, 0), i * height)
            lst._scroll_cursor_into_view()
            await pilot.pause(0)
            cells += _paint(lst)
        return cells

    if first:
        # Only on the first repetition: after it the pages are cached and this
        # would be a second warm sweep wearing the cold sweep's name.
        with _T("listing_cold_ms"):
            c = await sweep()
        note("listing_rows", lst.total)
        note("listing_cells", c)
    else:
        await sweep()
    with _T("listing_warm_ms"):
        await sweep()


async def phase_render(app, pilot, funcs, frames):
    """render_line throughput on a warm listing with the cursor moving every
    frame -- the steady state of holding down j."""
    lst = app.query_one(ListingView)
    app._active = "listing"
    app._show_active()
    lst.focus()
    await pilot.pause(0.02)
    n = 0
    with _T("render_ms"):
        for i in range(frames):
            lst.cursor = (i * 7) % max(lst.total, 1)
            for y in range(lst.size.height):
                n += len(lst.render_line(y).text)
    note("render_cells", n)


async def phase_decomp(app, pilot, funcs):
    """F5 out of the listing into the pseudocode, then read down it.

    Navigation is done OUTSIDE the timer -- phase_nav already charges for that,
    and counting it twice would let a nav win look like a decompiler win.
    """
    dv = app.query_one(DecompView)
    lst = app.query_one(ListingView)
    ok = lines = 0
    spent = 0.0
    for fn in funcs:
        app._open_function(fn.addr, fn.name)
        if not await _wait(pilot, lambda fn=fn: app._cur is not None
                           and app._cur.ea == fn.addr and lst.total > 0
                           and lst._cursor_ea() == fn.addr, 120):
            fail(f"decomp:open:{fn.name}")
            continue
        app._active = "listing"
        app._show_active()
        lst.focus()
        await pilot.pause(0.02)
        t0 = time.perf_counter()
        app.action_toggle_view()
        if not await _wait(pilot, lambda fn=fn: dv.loaded_ea == fn.addr, 60):
            fail(f"decomp:{fn.name}")
            continue
        app._active = "decomp"
        app._show_active()
        await pilot.pause(0)
        for _ in range(4):                  # scroll the body, painting each page
            _paint(dv)
            dv.cursor = min(max(dv.total - 1, 0), dv.cursor + dv.size.height)
            dv._scroll_cursor_into_view()
            await pilot.pause(0)
        spent += (time.perf_counter() - t0) * 1000
        ok += 1
        lines += dv.total
    TIMES.setdefault(PREFIX + "decomp_ms", []).append(spent)
    app._active = "listing"
    app._show_active()
    note("decomp_ok", ok)
    note("decomp_lines", lines)


async def phase_graph(app, pilot, funcs):
    gv = app.query_one(GraphView)
    ok = blocks = 0
    spent = 0.0
    # Space on a function you have not graphed before is the case that hurts:
    # it pays the flowchart tool, a heads walk over the function's extent, and
    # the layout. Program caches flowcharts per function, so without this the
    # phase would time a dict lookup.
    app.program._flowcharts.clear()
    for fn in funcs:
        app._open_function(fn.addr, fn.name)
        if not await _wait(pilot, lambda fn=fn: app._cur is not None
                           and app._cur.ea == fn.addr, 120):
            fail(f"graph:open:{fn.name}")
            continue
        app._graph_sticky = True
        gv.fc = None
        gv.lay = None
        t0 = time.perf_counter()
        app._load_graph(fn.addr, fn.addr)
        got = await _wait(pilot, lambda fn=fn: gv.fc is not None
                          and gv.fc.func_ea == fn.addr and gv.lay is not None, 60)
        app._graph_sticky = False
        if not got:
            fail(f"graph:{fn.name}")
            continue
        app._active = "graph"
        app._show_active()
        await pilot.pause(0)
        _paint(gv)
        for _ in range(3):                  # pan the canvas, painting each frame
            gv.scroll_to(y=min(gv.lay.height, round(gv.scroll_offset.y)
                               + gv.size.height), animate=False)
            await pilot.pause(0)
            _paint(gv)
        spent += (time.perf_counter() - t0) * 1000
        ok += 1
        blocks += len(gv.lay.nodes)
    TIMES.setdefault(PREFIX + "graph_ms", []).append(spent)
    app._active = "listing"
    app._show_active()
    note("graph_ok", ok)
    note("graph_blocks", blocks)


async def phase_split(app, pilot, funcs):
    """`s` — listing and pseudocode side by side, cursor-linked.

    The link is driven by ``decomp_map``: for every pseudocode line, the set of
    instructions the decompiler attributes to it. That is the whole cost of the
    feature and the bench did not cover it at all.
    """
    dv = app.query_one(DecompView)
    lst = app.query_one(ListingView)
    ok = mapped = 0
    spent = 0.0
    for fn in funcs:
        app._open_function(fn.addr, fn.name)
        if not await _wait(pilot, lambda fn=fn: app._cur is not None
                           and app._cur.ea == fn.addr and lst.total > 0
                           and lst._cursor_ea() == fn.addr, 120):
            fail(f"split:open:{fn.name}")
            continue
        app._active = "listing"
        app._show_active()
        lst.focus()
        await pilot.pause(0.02)
        t0 = time.perf_counter()
        app.action_toggle_split()
        got = await _wait(pilot, lambda fn=fn: app._split and lst.display
                          and dv.display and dv.loaded_ea == fn.addr, 120)
        # The region map is what the split view is FOR; wait for it, not just
        # for two panes to appear.
        m = app.program.decomp_map(fn.addr)
        spent += (time.perf_counter() - t0) * 1000
        if not got:
            fail(f"split:{fn.name}")
        else:
            ok += 1
            mapped += sum(1 for eas in m if eas)
        _paint(lst)
        _paint(dv)
        if app._split:
            app.action_toggle_split()
            await _wait(pilot, lambda: not app._split, 30)
    TIMES.setdefault(PREFIX + "split_ms", []).append(spent)
    app._active = "listing"
    app._show_active()
    note("split_ok", ok)
    note("split_mapped_lines", mapped)


async def phase_rename(app, pilot, funcs):
    """Rename a function and get the listing back — the commonest RE operation.

    A rename invalidates cached names everywhere, and the question this measures
    is what the *listing* then costs: the cursor has to land back on the same row
    with the new name showing. Renames are undone afterwards so the .i64 the
    bench stages from is never left edited.
    """
    lst = app.query_one(ListingView)
    ed = app.program
    done = 0
    spent = 0.0
    tag = f"_bench_{os.getpid()}"
    for k, fn in enumerate(funcs[:6]):
        app._open_function(fn.addr, fn.name)
        if not await _wait(pilot, lambda fn=fn: app._cur is not None
                           and app._cur.ea == fn.addr and lst.total > 0
                           and lst._cursor_ea() == fn.addr, 120):
            fail(f"rename:open:{fn.name}")
            continue
        row = lst.cursor
        new = f"{tag}_{k}"
        t0 = time.perf_counter()
        try:
            ed.client.call("rename",
                           batch={"func": [{"addr": hex(fn.addr), "name": new}]})
        except Exception:  # noqa: BLE001
            fail(f"rename:call:{fn.name}")
            continue
        ed.bump_names()
        # What the app does next: reopen the listing where it was and paint it.
        lm = ed.listing(fn.addr)
        idx = max(lm.ensure_ea(fn.addr), 0) if lm is not None else 0
        lst.load(lm, new, cursor=idx, scroll_y=max(idx - 6, 0))
        got = await _wait(pilot, lambda fn=fn: lst.total > 0
                          and lst._cursor_ea() == fn.addr, 120)
        await pilot.pause(0)
        _paint(lst)
        spent += (time.perf_counter() - t0) * 1000
        # The name lives on the `proc` banner row, a few rows above the code row
        # ensure_ea lands on (banner rows aren't address-indexed), so look at the
        # window rather than the single row.
        shows_new = False
        if lm is not None:
            for k2 in range(max(idx - 4, 0), idx + 2):
                h = lm.get(k2)
                if h is not None and (new in (h.text or "")
                                      or new == (h.name or "")):
                    shows_new = True
                    break
        if not (got and idx == row and shows_new):
            fail(f"rename:{fn.name}:row={idx}/{row} shows_new={shows_new}")
        else:
            done += 1
        try:                      # put it back, whatever happened above
            ed.client.call("rename",
                           batch={"func": [{"addr": hex(fn.addr),
                                            "name": fn.name}]})
            ed.bump_names()
        except Exception:  # noqa: BLE001
            fail(f"rename:undo:{fn.name}")
    TIMES.setdefault(PREFIX + "rename_ms", []).append(spent)
    note("rename_ok", done)


async def phase_search(app, pilot, funcs, terms):
    """The real incremental-search path: search_begin, then one search_update per
    typed character, exactly as the Input's on_changed drives it."""
    lst = app.query_one(ListingView)
    app._active = "listing"
    app._show_active()
    lst.focus()
    await pilot.pause(0.02)
    hits = 0
    with _T("search_ms"):
        for term in terms:
            lst.search_begin(1)
            for i in range(1, len(term) + 1):
                lst.search_update(term[:i])
                if not await _wait(pilot, lambda: bool(lst._matches), 180):
                    fail(f"search:{term}")
            hits += len(lst._matches)
            lst.search_cancel()
            await pilot.pause(0)
    note("search_hits", hits)


async def phase_index(app, pilot, funcs):
    from idatui.app import _fuzzy  # noqa: PLC0415
    prog = app.program
    with _T("index_ms"):
        for filt in (None, "sub_*", "*a*"):
            prog.functions(filt).load_all()
    n = len(prog.functions())
    note("index_n", n)
    if n == 0:
        fail("index:empty")
    names = [f.name for f in prog.functions().all_loaded()]
    with _T("palette_ms"):
        tot = 0
        for q in ("m", "mn", "sub", "prnt", "ab", "str"):
            for nm in names:
                tot += 1 if _fuzzy(nm, q) else 0
    note("palette_hits", tot)


async def phase_hex(app, pilot, funcs, frames):
    hx = app.query_one(HexView)
    fn = funcs[0]
    app._open_function(fn.addr, fn.name)
    await _wait(pilot, lambda: app._cur is not None and app._cur.ea == fn.addr, 60)
    app.action_hex()
    if not await _wait(pilot, lambda: app._active == "hex" and hx.model is not None, 60):
        fail("hex:open")
        return
    await pilot.pause(0.05)
    height = max(hx.size.height, 1)
    with _T("hex_ms"):
        for i in range(frames):
            row = (i * height) % max(hx.total, 1)
            hx.cursor = row * 16
            hx.model.ensure(row, height)
            hx._apply_scroll(row)
            await pilot.pause(0)
            _paint(hx)
    app._active = "listing"
    app._show_active()


def phase_pure_graph(corpus):
    from idatui import graph as G  # noqa: PLC0415
    with open(corpus) as fh:
        recs = json.load(fh)
    n = 0
    with _T("pure_graph_ms"):  # called three times; the median is taken
        for rec in recs:
            texts = {b["id"]: list(b["lines"]) for b in rec["blocks"]}
            blocks = [G.Block(id=b["id"], start=b["start"], end=b["end"],
                              succs=[(d, k) for d, k in b["succs"]])
                      for b in rec["blocks"]]

            def sizer(b, texts=texts):
                lines = texts[b.id]
                label = f"loc_{b.start:X}"
                widest = max([len(label) + 4] + [len(ln) for ln in lines] or [4])
                return (widest + 4, max(len(lines), 1) + 2)

            lay = G.layout(blocks, sizer)
            n += len(lay.nodes)
    NOTES["pure_graph_nodes"] = n
    if n == 0:
        FAILS.append("pure_graph:empty")


# --------------------------------------------------------------------------- #
def stage(binary):
    """A temp-dir copy of ``binary`` seeded from its pre-analysed database.

    Prefers ``<bin>.pristine.i64`` (what tests/_fixtures builds) and falls back
    to a tracked ``<bin>.i64``. Never touches either.
    """
    d = tempfile.mkdtemp(prefix="idatui-bench-")
    tgt = os.path.join(d, os.path.basename(binary))
    shutil.copy2(binary, tgt)
    for cand in (binary + ".pristine.i64", binary + ".i64"):
        if os.path.exists(cand):
            shutil.copy2(cand, tgt + ".i64")
            break
    else:
        raise SystemExit(f"{binary}: no pre-analysed .i64 -- bench would time "
                         f"auto-analysis, not ida-tui")
    return d, tgt


async def run_target(binary, reps, nfuncs, pages, frames, terms, skip=0):
    d, target = stage(binary)
    try:
        app = IdaTui(open_path=target, keepalive=False)
        t0 = time.perf_counter()
        async with app.run_test(size=SIZE) as pilot:
            # Usable = index complete AND the loading overlay handed the screen
            # back (auto-land fired). Benching with the modal still up silently
            # disables every priority binding, so nothing a phase presses lands.
            if not await _wait(pilot, lambda: app._func_index is not None
                               and app._func_index.complete, 600):
                fail("boot:index")
                return
            await _wait(pilot, lambda: app._loading_screen is None
                        and len(app.screen_stack) == 1, 120)
            TIMES.setdefault(PREFIX + "boot_ms", []).append(
                (time.perf_counter() - t0) * 1000)
            idx = app._func_index
            idx.load_all()
            allf = idx.all_loaded()
            note("funcs", len(allf))
            # Largest by size, ties by address, skipping any function the graph
            # view refuses (>GRAPH_MAX_BLOCKS) -- otherwise the graph phase sits
            # in a timeout instead of doing work.
            # ``skip`` drops the very largest: on a big binary the top few are
            # multi-thousand-line monsters whose Hex-Rays cost is 30s of pure
            # backend time, which would swamp every phase we can actually make
            # faster. A band just below the top is still large and realistic.
            big = []
            for f in sorted(allf, key=lambda f: (-f.size, f.addr))[skip:]:
                if len(big) >= nfuncs:
                    break
                try:
                    # The raw tool, NOT Program.flowchart: the latter caches, and
                    # picking the fixtures through it left phase_graph measuring
                    # cache hits instead of the work a user waits for.
                    payload = app.program.client.call("flowchart",
                                                      addr=hex(f.addr))
                except Exception:  # noqa: BLE001
                    continue
                if not isinstance(payload, dict) or payload.get("error"):
                    continue
                nblocks = len(payload.get("blocks") or [])
                if not nblocks or nblocks > IdaTui.GRAPH_MAX_BLOCKS:
                    continue
                big.append(f)
            if len(big) < nfuncs:
                fail(f"fixture:only {len(big)}/{nfuncs} functions")
            big = sorted(big, key=lambda f: f.addr)
            note("fixed_set", [f.name for f in big])

            # Cold, once: the listing model starts empty and nav is the thing
            # that has to walk it. Repeating it would only measure a warm cache.
            await phase_nav(app, pilot, big)

            # Phases split two ways. The repeatable ones (paint throughput, and
            # the graph, which clears its own cache) are run every repetition so
            # the median settles. The cold-sensitive ones -- decompile, search
            # and the function index all cache their answer, and re-running them
            # would report a dict lookup under the name of the thing a user
            # waits for -- are run ONCE.
            for rep in range(reps):
                first = rep == 0
                await phase_listing(app, pilot, big, pages, first)
                await phase_render(app, pilot, big, frames)
                if first:
                    await phase_decomp(app, pilot, big)
                await phase_graph(app, pilot, big)
                if first:
                    await phase_rename(app, pilot, big)
                    await phase_split(app, pilot, big)
                    await phase_search(app, pilot, big, terms)
                    await phase_index(app, pilot, big)
                await phase_hex(app, pilot, big, frames // 5)
            app.exit()
    finally:
        shutil.rmtree(d, ignore_errors=True)


def main():
    global PREFIX
    ap = argparse.ArgumentParser()
    ap.add_argument("--corpus", default=os.path.join(HERE, "cfg-corpus.json"))
    ap.add_argument("--only", default="", help="only this target tag")
    ap.add_argument("--reps", type=int, default=0, help="override reps")
    a = ap.parse_args()

    plan = [
        # tag,   binary,        reps, nfuncs, pages, frames, terms, skip
        ("sm_", "targets/echo", 2, 12, 40, 300, ("mov", "call", "lea"), 0),
        # Two reps on the big target too: one sample of a 2.5s Hex-Rays phase
        # swings by 10%, which was drowning changes worth more than that.
        # (phase_nav stays single -- it is a COLD walk by definition.)
        ("lg_", "targets/bash", 2, 12, 60, 300, ("mov", "call"), 120),
    ]
    if a.only:
        plan = [p for p in plan if p[0].startswith(a.only)]

    async def go():
        global PREFIX
        for tag, binary, reps, nfuncs, pages, frames, terms, skip in plan:
            PREFIX = tag
            await run_target(os.path.join(ROOT, binary), a.reps or reps,
                             nfuncs, pages, frames, terms, skip)
        PREFIX = ""

    asyncio.run(go())
    if os.path.exists(a.corpus):
        # Three times, for the median. This phase is pure CPU with no I/O and
        # was the single jumpiest number in the suite (239 <-> 524 with
        # identical code) purely from being descheduled -- which is noise the
        # primary metric was carrying for no reason.
        for _ in range(3):
            phase_pure_graph(a.corpus)
    else:
        FAILS.append("pure_graph:no corpus")

    out = {}
    for k, v in TIMES.items():
        out[k] = statistics.median(v) if len(v) > 2 else min(v)
    total = sum(out.values())
    print()
    for k in sorted(out):
        print(f"METRIC {k}={out[k]:.1f}")
    print(f"METRIC total_ms={total:.1f}")
    print(f"METRIC fails={len(FAILS)}")
    print("NOTES " + json.dumps(NOTES)[:3000])
    if FAILS:
        print("FAILS " + "; ".join(FAILS[:20]))
        return 1
    return 0


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