#!/usr/bin/env python3
"""Observer side of the black-box glass-to-glass instrument.

    client_probe.py --window latenz_app [--samples 30] [--interval 0.05]

Reads ONE timecode -- the client's rendering of the strip that g2g_probe draws
inside the streamed desktop -- and subtracts it from this machine's
CLOCK_MONOTONIC. The streamed desktop is headless (a container, a VM, another
box), so its live copy is not in this machine's framebuffer and the
two-timecode instrument in g2g_measure.py does not apply: that one locates two
green boxes in one root screenshot and pairs them, and here there is only ever
one. Use this module for a containerised source, g2g_measure.py for a source
that is this machine's own X11 root.

Because the source-side copy is no longer in the picture, the probe's own
draw->compose no longer cancels between two timecodes. The number this
produces therefore INCLUDES the streamed desktop's compositor, and is larger
than -- and not comparable to -- the two-timecode figures in the viewer's
historical run records. That is the honest end-to-end quantity; say so
wherever it is published.

Everything that reads pixels is imported from g2g_measure rather than
reimplemented: Grabber (per-rect XGetImage, because grabbing the whole root
stalls the compositor ~200 ms and the instrument then reads its own stalls
back as latency), green_boxes (the hollow-frame band pairing) and decode.
"""
import argparse
import json
import os
import pathlib
import re
import sys
import time

sys.path.insert(0, str(pathlib.Path(__file__).resolve().parent))
import g2g_measure as g  # noqa: E402

try:
    import numpy as np
except ImportError:
    np = None

from Xlib import display as xdisplay  # noqa: E402

DISPLAY = os.environ.get("G2G_DISPLAY", os.environ.get("DISPLAY", ":0"))

# The aspect is a sanity bound, NOT the anti-stretch test. Measured on real
# streams: the same strip reads 10.00 over H.264 at scale 0.667 and 11.31 over
# H.265, against the source's 9.25. Chroma subsampling and deblocking erode the
# green border's top and bottom rows while leaving its width intact, so the
# measured height runs short by an amount that depends on the codec -- up to
# 22%, which is larger than the 11% a genuine 16:9-into-16:10 stretch produces.
# Aspect therefore cannot separate erosion from stretching, and a tight bound
# here only rejects good measurements. Erosion does not disturb decode: it eats
# both edges, so the block centres stay centred.
#
# The real anti-stretch test is `expect_scale` in locate_timecode: the strip's
# width-derived scale must match the scale the video is known to be displayed
# at. That uses independent information instead of the strip's own geometry.
ASPECT_TOL = 0.50
SCALE_MATCH_TOL = 0.10
SCALE_MIN, SCALE_MAX = 0.25, 1.5


class NoTimecode(Exception):
    """The strip could not be located, or was located ambiguously.

    Never guessed around: a mislocated pattern reports a frozen or wrong
    timecode, which looks exactly like a stalled stream and has invalidated a
    campaign in this repo before.
    """


def list_windows(display=None):
    """Every window with a WM_NAME, in root coordinates.

    Walks children and grandchildren like the existing g2g_latency test does,
    but translates to root coordinates -- the client window is reparented by
    the window manager, so its own x/y are relative to the frame, not the root.
    """
    dpy = xdisplay.Display(display or DISPLAY)
    root = dpy.screen().root
    out = []
    for w in root.query_tree().children:
        try:
            children = [w] + list(w.query_tree().children)
        except Exception:
            # QueryTree on a window that died between the root scan and here
            # raises BadWindow and takes the whole enumeration with it.
            continue
        for s in children:
            try:
                nm = s.get_wm_name()
            except Exception:
                continue
            if not nm:
                continue
            try:
                geom = s.get_geometry()
                t = s.translate_coords(root, 0, 0)
            except Exception:
                # A window can die between query_tree and these two round
                # trips; X answers BadWindow and the whole scan dies with it.
                # Sampling in a loop hits this eventually -- once per ~2 min
                # here, mid-run, which reads as an instrument failure.
                continue
            # translate_coords gives the offset of root's origin in s's frame,
            # so the window's root position is its negation.
            out.append({"name": nm, "id": s.id, "x": -t.x, "y": -t.y,
                        "w": geom.width, "h": geom.height})
    return out


def find_window_near(name, x, y, display=None):
    """The window called `name` whose origin is nearest (x, y).

    With one session per codec there are two identically named viewers on the
    screen at once, and find_window refuses ambiguity by design. The rig knows
    where it put each one, so position is the disambiguator -- and it is
    checked, not assumed: a window further than half a cell from where the rig
    placed it is reported rather than silently measured, because measuring the
    wrong viewer is exactly the failure that produced a retracted result here.
    """
    hits = [w for w in list_windows(display) if w["name"] == name]
    if not hits:
        raise NoTimecode(f"no window named {name!r}")
    best = min(hits, key=lambda w: (w["x"] - x) ** 2 + (w["y"] - y) ** 2)
    if abs(best["x"] - x) > best["w"] // 2 or abs(best["y"] - y) > best["h"] // 2:
        raise NoTimecode(
            f"the nearest {name!r} is at {best['x']},{best['y']}, more than half "
            f"a window from the {x},{y} the rig placed it at -- refusing to "
            f"guess which viewer this is")
    return best


def find_window(name, display=None, exact=True):
    """The one window called `name`, matched exactly by default.

    Exact, and not a regex search, from a bug that cost an afternoon: the
    Flutter viewer owns TWO named toplevels, `latenz_app` and
    `com.example.latenz_app`. A substring match hits whichever the tree walk
    reaches first, and that order is not stable across launches -- so the rig
    would place one window while the observer measured the other, and the
    symptom was an intermittent "no timecode found" against a stream that was
    demonstrably decoding frames. Ambiguity is now an error, and the error
    lists the candidates.
    """
    wins = list_windows(display)
    if exact:
        hits = [w for w in wins if w["name"] == name]
    else:
        pat = re.compile(name)
        hits = [w for w in wins if pat.search(w["name"])]
    if len(hits) == 1:
        return hits[0]
    if not hits:
        raise NoTimecode(
            f"no window named {name!r} on {display or DISPLAY}. Present: "
            f"{sorted({w['name'] for w in wins})}")
    raise NoTimecode(
        f"{len(hits)} windows match {name!r}: "
        f"{[(w['name'], w['w'], w['h'], w['x'], w['y']) for w in hits]}. "
        f"Name exactly one of them -- measuring the wrong one reads as a dead "
        f"stream.")


def _green_mask_boxes(img, min_w=120):
    """green_boxes with a numpy mask when numpy is available.

    green_boxes as written is a per-pixel Python loop and is the only expensive
    step in a sample; the mask is the part worth vectorising. The run grouping
    below is left in Python because there are only a handful of runs.
    """
    if np is None:
        return g.green_boxes(img, min_w=min_w)
    a = np.asarray(img, dtype=np.int16)
    r, gr, b = a[:, :, 0], a[:, :, 1], a[:, :, 2]
    mask = (gr > 110) & (r < gr - 60) & (b < gr - 60)
    if not mask.any():
        return []
    pad = np.zeros((mask.shape[0], mask.shape[1] + 2), dtype=bool)
    pad[:, 1:-1] = mask
    d = np.diff(pad.astype(np.int8), axis=1)
    runs = []
    for y in np.flatnonzero(mask.any(axis=1)):
        starts = np.flatnonzero(d[y] == 1)
        ends = np.flatnonzero(d[y] == -1)
        for x0, x1 in zip(starts, ends):
            if x1 - x0 >= min_w:
                runs.append((int(y), int(x0), int(x1)))

    groups = []
    for y, x0, x1 in runs:
        for gx in groups:
            if abs(gx[0] - x0) < 0.1 * (x1 - x0) and abs(gx[2] - x1) < 0.1 * (x1 - x0):
                gx[0], gx[2], gx[3] = min(gx[0], x0), max(gx[2], x1), y
                break
        else:
            groups.append([x0, y, x1, y])

    out = []
    for x0, y0, x1, y1 in groups:
        bw, bh = x1 - x0, y1 - y0 + 1
        if bh < 8:
            continue
        if 0.6 * g.ASPECT < bw / bh < 1.6 * g.ASPECT:
            out.append((x0, y0, bw, bh))
    out.sort(key=lambda b: (b[1], b[0]))
    return out


def locate_timecode(grabber, win, expect_scale=None):
    """Find the single strip inside the client window, in root coordinates.

    Searches the window rectangle only -- never the root. Fails loudly on zero
    or several candidates, on a stretched aspect, or on a strip touching the
    window edge (which means it may be cropped, and a cropped strip decodes to
    a wrong number instead of failing).
    """
    img = grabber.rect(win["x"], win["y"], win["w"], win["h"])
    boxes = _green_mask_boxes(img)
    if len(boxes) != 1:
        raise NoTimecode(
            f"expected exactly one timecode inside {win['name']}, found "
            f"{len(boxes)}: {boxes}. Zero usually means the stream is not "
            f"running, the probe is not started, or the client is scaling the "
            f"strip below the aspect filter; more than one means something "
            f"else green is in the window.")
    bx, by, bw, bh = boxes[0]
    aspect = bw / bh
    if abs(aspect - g.ASPECT) / g.ASPECT > ASPECT_TOL:
        raise NoTimecode(
            f"timecode aspect {aspect:.2f} is nowhere near the source's "
            f"{g.ASPECT:.2f} -- this is probably not the strip.")
    scale = bw / g.PW
    if expect_scale and abs(scale / expect_scale - 1.0) > SCALE_MATCH_TOL:
        raise NoTimecode(
            f"timecode scale {scale:.3f} does not match the {expect_scale:.3f} "
            f"the video is displayed at -- the client is scaling the picture "
            f"non-uniformly, so decoded block positions would be skewed.")
    if not (SCALE_MIN <= scale <= SCALE_MAX):
        raise NoTimecode(f"timecode scale {scale:.3f} outside "
                         f"[{SCALE_MIN}, {SCALE_MAX}]")
    if bx <= 0 or by <= 0 or bx + bw >= win["w"] or by + bh >= win["h"]:
        raise NoTimecode(
            f"timecode touches the window edge ({bx},{by},{bw},{bh} in "
            f"{win['w']}x{win['h']}) -- it may be cropped. Increase the "
            f"probe's --inset or enlarge the client window.")
    return {"x": win["x"] + bx, "y": win["y"] + by, "w": bw, "h": bh,
            "scale": scale, "aspect": aspect}


def sample_latency(grabber, box, n, interval):
    """n samples of (latency_ms, grab_us, timecode).

    The timestamp is the midpoint of the framebuffer read, not its start: the
    read is not instantaneous and charging its whole duration to either end
    biases every sample in one direction.
    """
    out = []
    for _ in range(n):
        # Discarded 16x16 read first. Measured: back to back, a strip grab costs
        # 0.26 ms; with a sleep between samples, one in five costs 10-16 ms,
        # because the first read after an idle gap pays a resync on the region
        # the compositor has since handed to the GPU. The warm-up pays it, so
        # the timed read does not, and the timestamp's uncertainty stays at the
        # sub-millisecond level. (Latency itself does not move: p50 was 52, 60,
        # 52 ms at 500 ms, 0 ms and 50 ms sampling intervals -- the instrument
        # is not the thing being measured. That control is worth re-running
        # whenever this path changes.)
        grabber.rect(box["x"], box["y"], 16, 16)
        t0 = time.monotonic()
        img = grabber.rect(box["x"], box["y"], box["w"], box["h"])
        t1 = time.monotonic()
        v = g.decode(img, 0, 0, box["w"], box["h"])
        if v is None:
            out.append((None, (t1 - t0) * 1e6, None))
        else:
            mid_ms = int(((t0 + t1) / 2) * 1000)
            out.append(((mid_ms - v) % 4096, (t1 - t0) * 1e6, v))
        if interval:
            time.sleep(interval)
    return out


def sample_unique_fps(grabber, box, seconds, hz=200):
    """Distinct decoded timecodes per second: smoothness measured at the glass.

    Needs no cooperation from any component, which is the point -- it is the
    one frame-rate figure that means the same thing on any stack.
    """
    period = 1.0 / hz
    seen, end, total = set(), time.monotonic() + seconds, 0
    nxt = time.monotonic()
    while time.monotonic() < end:
        img = grabber.rect(box["x"], box["y"], box["w"], box["h"])
        v = g.decode(img, 0, 0, box["w"], box["h"])
        total += 1
        if v is not None:
            seen.add(v)
        nxt += period
        d = nxt - time.monotonic()
        if d > 0:
            time.sleep(d)
        else:
            nxt = time.monotonic()
    return {"unique_fps": len(seen) / seconds, "polls": total,
            "poll_hz": total / seconds}


def percentile(vals, p):
    if not vals:
        return None
    s = sorted(vals)
    k = min(len(s) - 1, max(0, int(round((p / 100.0) * (len(s) - 1)))))
    return s[k]


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--window", default="latenz_app",
                    help="regex matched against WM_NAME")
    ap.add_argument("--samples", type=int, default=30)
    ap.add_argument("--interval", type=float, default=0.05)
    ap.add_argument("--unique-fps-seconds", type=float, default=0.0)
    ap.add_argument("--json", action="store_true")
    a = ap.parse_args()

    win = find_window(a.window)
    grabber = g.Grabber()
    box = locate_timecode(grabber, win)
    rows = sample_latency(grabber, box, a.samples, a.interval)

    lat = [r[0] for r in rows if r[0] is not None]
    grabs = [r[1] for r in rows]
    codes = [r[2] for r in rows if r[2] is not None]
    distinct = sum(1 for i in range(1, len(codes)) if codes[i] != codes[i - 1])
    res = {
        "window": win, "box": box,
        "n": len(rows), "decoded": len(lat),
        "p50_ms": percentile(lat, 50), "p90_ms": percentile(lat, 90),
        "in_range_frac": sum(1 for v in lat if 0 < v < 1000) / len(lat) if lat else 0.0,
        "distinct_frac": distinct / (len(codes) - 1) if len(codes) > 1 else 0.0,
        "grab_us_p50": percentile(grabs, 50),
        "grab_us_p90": percentile(grabs, 90),
        "grab_us_max": max(grabs) if grabs else None,
    }
    if a.unique_fps_seconds:
        res.update(sample_unique_fps(grabber, box, a.unique_fps_seconds))

    if a.json:
        print(json.dumps(res, indent=2))
    else:
        print(f"window   {win['name']} {win['w']}x{win['h']}+{win['x']}+{win['y']}")
        print(f"timecode {box['w']}x{box['h']}+{box['x']}+{box['y']} "
              f"scale {box['scale']:.3f} aspect {box['aspect']:.2f}")
        print(f"p50      {res['p50_ms']} ms   p90 {res['p90_ms']} ms")
        print(f"in range {res['in_range_frac']:.0%}   distinct "
              f"{res['distinct_frac']:.0%}   grab p50 {res['grab_us_p50']:.0f} / "
              f"p90 {res['grab_us_p90']:.0f} us")
        if a.unique_fps_seconds:
            print(f"unique   {res['unique_fps']:.1f} fps "
                  f"(polled at {res['poll_hz']:.0f} Hz)")
    return 0


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