"""G0 probe — đo latency từng plugin + jitter round-trip ≥10k block.

Dùng ĐÚNG pipeline thật (FxRealtimeSession → SHM → fx_vst_bridge → SHM),
không sửa app code. WS browser → Python không đo được từ host; phần đó lấy
bằng chứng từ dbg.log (fx_play_first calib + fx_stats u/s) — probe này đo
Python pump (batching) + SHM + bridge + return.

Chạy:  python native_bridge/debug/g0_latency_probe.py [--blocks 10000]
"""
import argparse
import atexit
import ctypes
import os
import statistics
import sys
import time

sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..")))

# Windows Sleep granularity mặc định 15.6ms → probe pacing ~22ms/batch ≠
# realtime 21.33ms → số liệu round-trip lệch. Bridge C++ đã gọi
# timeBeginPeriod(1) (RealtimeFxLoop.cpp:160); probe phải làm tương tự để
# pacing đúng realtime.
try:
    _winmm = ctypes.windll.winmm
    _winmm.timeBeginPeriod(1)
    atexit.register(_winmm.timeEndPeriod, 1)
except Exception:
    pass

import numpy as np  # noqa: E402

from app.core.fx_realtime import FxRealtimeSession, FXRT_BLOCK  # noqa: E402

SR = 48000
PERIOD = FXRT_BLOCK / SR  # 5.333ms/block; batch 4 = 21.333ms

PURECOMP = r"C:\Program Files\Common Files\VST3\Sonible\purecomp.vst3"
OZONE = r"C:\Program Files\Common Files\VST3\iZotope\Ozone 11.vst3"

CONFIGS = [
    ("empty (passthrough)", []),
    ("pureComp", [{"type": "vst3", "path": PURECOMP, "bypass": False, "preset_b64": ""}]),
    ("Ozone 11", [{"type": "vst3", "path": OZONE, "bypass": False, "preset_b64": ""}]),
]


def make_block():
    t = np.arange(FXRT_BLOCK, dtype=np.float32) / SR
    s = (np.sin(2 * np.pi * 440 * t) * 0.1).astype(np.float32)
    inter = np.empty(FXRT_BLOCK * 2, dtype=np.float32)
    inter[0::2] = s
    inter[1::2] = s
    return inter.tobytes()


def run(name, chain, n_blocks):
    log_path = os.path.join(os.environ.get("TEMP", "."), "g0_probe_%d.log" % os.getpid())
    os.environ["FXRT_BRIDGE_LOG"] = log_path
    sess = FxRealtimeSession(chain, SR)
    stats = {}
    try:
        sess.wait_ready(30)
        time.sleep(1.0)  # chain ready + latency report
        stats["plugin_latency"] = dict(sess.get_latencies())
        blk = make_block()
        send_t = []   # block idx -> send time (call time)
        recv_t = []   # (block idx, recv time) — FIFO: out block j <=> in block j
        next_idx = 0
        batch_dur = PERIOD * 4
        batches = n_blocks // 4
        pump_t0 = time.perf_counter()
        for b in range(batches):
            t0 = time.perf_counter()
            for _ in range(4):
                send_t.append(time.perf_counter())
                sess.write_input(blk)
                next_idx += 1
            sess.flush_input_stale(0.001)
            outs = sess.read_output()
            base = len(recv_t)
            for k in range(len(outs)):
                recv_t.append((base + k, time.perf_counter()))
            dt = time.perf_counter() - t0
            if dt < batch_dur:
                time.sleep(batch_dur - dt)
        # Drain phần còn lại (bridge xử lí nốt backlog) — nếu bridge theo kịp
        # input, recv đủ = sent ngay sau pump; thiếu kéo dài = drift thật.
        drain_t0 = time.perf_counter()
        while len(recv_t) < len(send_t) and time.perf_counter() - drain_t0 < 2.0:
            outs = sess.read_output()
            base = len(recv_t)
            for k in range(len(outs)):
                recv_t.append((base + k, time.perf_counter()))
            if not outs:
                time.sleep(0.002)
        pump_t1 = time.perf_counter()
        stats["pump_rate"] = len(recv_t) / (pump_t1 - pump_t0)
        stats["drained"] = len(recv_t) == len(send_t)
        # round-trip per block
        rt = []
        for j, tr in recv_t:
            if j < len(send_t):
                rt.append((tr - send_t[j]) * 1000.0)
        # inter-arrival deviation vs ideal period
        recv_abs = [tr for _, tr in recv_t]
        ia = []
        for k in range(1, len(recv_abs)):
            ia.append((recv_abs[k] - recv_abs[k - 1]) * 1000.0 - PERIOD * 1000.0)
        stats["roundtrip_ms"] = {
            "n": len(rt), "mean": statistics.mean(rt), "sd": statistics.pstdev(rt),
            "min": min(rt), "max": max(rt),
            "p50": sorted(rt)[len(rt) // 2],
            "p99": sorted(rt)[int(len(rt) * 0.99) - 1],
        }
        stats["interarrival_dev_ms"] = {
            "n": len(ia), "mean": statistics.mean(ia), "sd": statistics.pstdev(ia),
            "min": min(ia), "max": max(ia),
        }
        stats["sent"] = len(send_t)
        stats["recv"] = len(recv_t)
        # bridge FxRTPerf
        perf = []
        if os.path.exists(log_path):
            with open(log_path, "rb") as f:
                data = f.read().decode("utf-8", "replace")
            for line in data.splitlines():
                if "[FxRTPerf]" in line:
                    perf.append(line)
        if perf:
            stats["bridge_perf_last"] = perf[-1].split("] ", 1)[-1]
    finally:
        sess.close()
        try:
            os.remove(log_path)
        except OSError:
            pass
    return stats


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--blocks", type=int, default=10000)
    args = ap.parse_args()
    print("G0 probe — SR=%d block=%d n_blocks=%d (%.1fs/config)" % (SR, FXRT_BLOCK, args.blocks, args.blocks * PERIOD))
    results = {}
    for name, chain in CONFIGS:
        print("\n=== %s ===" % name, flush=True)
        st = run(name, chain, args.blocks)
        results[name] = st
        print("plugin_latency:", st.get("plugin_latency"))
        print("pump_rate=%.3f blk/s drained=%s (bridge theo kip input: rate ~187.5 va drained=True)" %
              (st.get("pump_rate", 0), st.get("drained")))
        rt = st.get("roundtrip_ms", {})
        print("roundtrip_ms n=%d mean=%.3f sd=%.3f min=%.3f p50=%.3f p99=%.3f max=%.3f" %
              (rt.get("n", 0), rt.get("mean", 0), rt.get("sd", 0), rt.get("min", 0),
               rt.get("p50", 0), rt.get("p99", 0), rt.get("max", 0)))
        ia = st.get("interarrival_dev_ms", {})
        print("interarrival_dev_ms n=%d mean=%.3f sd=%.3f min=%.3f max=%.3f" %
              (ia.get("n", 0), ia.get("mean", 0), ia.get("sd", 0), ia.get("min", 0), ia.get("max", 0)))
        print("bridge:", st.get("bridge_perf_last", "n/a"))
    print("\n=== JSON ===")
    print(json_dump(results))


def json_dump(results):
    import json
    return json.dumps(results, indent=1)


if __name__ == "__main__":
    main()