"""方案 R4 自动探针：无头 Edge / Chrome 里测浏览器推理延迟（部署形态：1 行特征 + K=8 条候选）。
  python -m ml.eval.r4_probe --model ml/runs/r4/traj_policy_r4.onnx --out ml/runs/r4.json
做法：本进程起一个带 COOP/COEP 头的静态服务（根目录 = 仓库），逐个配置用无头浏览器打开 webapp/r4_probe.html，
页面把 P50/P95/P99 与 20 ms 节拍最大间隔 POST 回来。配置：
  wasm·默认线程（与查看器同一默认 min(16, ⌊0.6×逻辑核⌋)，推理在 Worker）——部署默认口径
  wasm·4 线程（第四轮实测的旧默认，作对照）
  wasm·默认线程·不开 Worker（看推理是否卡主线程：tickGapMax）
  webgpu
判据：部署口径 P99 < 15 ms 即通过；wasm 不过而 WebGPU 过，也判通过并记 deploy_ep=webgpu（查看器用 &ep=webgpu）。
找不到浏览器 / 页面没回报 → pass=null（流水线警告继续，验收标红）。浏览器路径可用 R4_BROWSER 指定。
2026-10-18 第五轮新增（审查第四轮缺口 3：R4 需要人开浏览器看，导致 `all` 天生不完整）。
范围说明：R4 判据只看纯推理（随机权重只能代表推理开销）。
2026-10-18 第六轮（审查 B-1/B-2/B-3）：
  --viewer：在真查看器里跑 autotest（与 Node bench 同一批 30 道贴立柱题），记录整次规划 planMs（推理 + 前瞻 + 安全层）、
    节拍间隔与饥饿（丢步）次数。r4 阶段用随机模型跑一次，只作参考（随机策略的安全层路径不代表真模型）；
    bench 阶段用训练好的模型跑，写 ml/runs/r9_browser.json，判据 = 方案 R9：planMs P99 < 45 ms（60 ms 周期留 15 ms 余量）且饥饿节拍 < 1%。
  本地有 webapp/vendor/ort/（bash ml/pipeline.sh setup 会放好）就用本地 ort，不依赖 CDN（第六轮实测 jsdelivr 间歇不可达）。
  WebGPU 与查看器两种配置用有界面窗口打开（无头模式没有真 GPU 适配器，第六轮实测 backend not found），跑完自动关闭。"""
import argparse, http.server, json, os, shutil, socket, subprocess, sys, tempfile, threading, time, urllib.parse

ROOT = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
LIMIT_MS = 15.0
PLAN_LIMIT_MS = 45.0      # R9：整次规划 P99（60 ms 周期留 15 ms 给渲染与 tick）
STARVE_LIMIT = 0.01       # R9：饥饿节拍占比上限（只计运动中，冷启动不计）
TIME_LIMIT = 0.05         # R9：认证因墙钟截止判不过的占比上限——防止靠拒绝候选、刹停换来 planP99 达标（第八轮，viewer_ref 随机模型 time 占失败 70%）
# 2026-10-18 第八轮：R9 判稳态 planWarmP99（剔除每题首次冷启动规划）。冷启动发生在静止起步，只是起步延迟，单独报 planColdP99 不设闸；
# 否则 30 题 = 30 次冷启动，planP99 测的是题数而不是稳态。


def ort_local():
    # 2026-10-18 第八轮：ort 已入库（webapp/vendor/ort），默认只用本地；本地缺失直接报错，不再静默退回 CDN
    # （第七轮本机：npm pack 失败 → 退回 CDN → jsDelivr 不通 → 四种配置全是「加载 ort 失败」却显示跑完）。确要用 CDN 设 ORT_LOCAL=0。
    if os.environ.get('ORT_LOCAL') == '0': return False
    if not os.path.exists(os.path.join(ROOT, 'webapp', 'vendor', 'ort', 'ort.min.js')):
        sys.exit('webapp/vendor/ort/ort.min.js 不存在：git pull 拿入库的 ort，或 bash ml/prep/vendor_ort.sh；确要用 CDN 设 ORT_LOCAL=0')
    return True


def find_browser():
    if os.environ.get('R4_BROWSER'): return os.environ['R4_BROWSER']
    cands = []
    for base in (os.environ.get('ProgramFiles(x86)'), os.environ.get('ProgramFiles'), os.environ.get('LOCALAPPDATA')):
        if base:
            cands += [os.path.join(base, 'Microsoft', 'Edge', 'Application', 'msedge.exe'),
                      os.path.join(base, 'Google', 'Chrome', 'Application', 'chrome.exe')]
    for n in ('microsoft-edge', 'microsoft-edge-stable', 'google-chrome', 'google-chrome-stable', 'chromium', 'chromium-browser'):
        p = shutil.which(n)
        if p: cands.append(p)
    return next((c for c in cands if c and os.path.exists(c)), None)


class Srv:
    def __init__(self):
        self.result, self.ev = None, threading.Event()
        me = self

        class H(http.server.SimpleHTTPRequestHandler):
            extensions_map = {**http.server.SimpleHTTPRequestHandler.extensions_map, '.wasm': 'application/wasm', '.mjs': 'text/javascript'}

            def __init__(s, *a, **k): super().__init__(*a, directory=ROOT, **k)

            def end_headers(s):
                s.send_header('Cross-Origin-Opener-Policy', 'same-origin'); s.send_header('Cross-Origin-Embedder-Policy', 'require-corp')
                s.send_header('Cross-Origin-Resource-Policy', 'cross-origin'); s.send_header('Cache-Control', 'no-store')
                super().end_headers()

            def do_POST(s):
                n = int(s.headers.get('Content-Length') or 0)
                try: me.result = json.loads(s.rfile.read(n) or b'{}')
                except Exception as e: me.result = {'error': 'bad json: %s' % e}
                me.ev.set(); s.send_response(200); s.end_headers(); s.wfile.write(b'ok')

            def log_message(s, *a): pass

        sk = socket.socket(); sk.bind(('127.0.0.1', 0)); self.port = sk.getsockname()[1]; sk.close()
        self.httpd = http.server.ThreadingHTTPServer(('127.0.0.1', self.port), H)
        threading.Thread(target=self.httpd.serve_forever, daemon=True).start()


def kill_tree(p):
    if os.name == 'nt': subprocess.run(['taskkill', '/T', '/F', '/PID', str(p.pid)], capture_output=True)
    else:
        try: p.terminate(); p.wait(5)
        except Exception: p.kill()


def open_page(srv, browser, url, timeout_s, headed=False):
    srv.result = None; srv.ev.clear()
    prof = tempfile.mkdtemp(prefix='r4prof_')
    args = [browser] + ([] if headed else ['--headless=new']) + [f'--user-data-dir={prof}', '--no-first-run', '--no-default-browser-check',
            '--disable-extensions', '--enable-unsafe-webgpu', '--disable-background-timer-throttling', '--disable-renderer-backgrounding',
            '--window-size=1280,800', url]
    p = subprocess.Popen(args, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
    got = srv.ev.wait(timeout_s)
    kill_tree(p); shutil.rmtree(prof, ignore_errors=True)
    return srv.result if got else {'error': f'{timeout_s} s 内页面没有回报'}


def one(srv, browser, model_url, ep, threads, proxy, n, timeout_s):
    q = urllib.parse.urlencode({'model': model_url, 'ep': ep, 'threads': threads, 'proxy': proxy, 'n': n, **({'ortLocal': '1'} if ort_local() else {})})
    r = open_page(srv, browser, f'http://127.0.0.1:{srv.port}/webapp/r4_probe.html?{q}', timeout_s, headed=(ep == 'webgpu'))
    r = dict(r or {}, config=f'{ep}·threads={threads}·proxy={proxy}')
    print(json.dumps(r, ensure_ascii=False), flush=True)
    return r


def viewer(srv, browser, model_url, n, probe):
    """真查看器 autotest：整次规划耗时 / 节拍 / 饥饿（R9 口径）"""
    q = urllib.parse.urlencode({'policy': model_url, 'autotest': n, **({'probe': '1'} if probe else {}), **({'ortLocal': '1'} if ort_local() else {})})
    r = open_page(srv, browser, f'http://127.0.0.1:{srv.port}/webapp/index.html?{q}', 120 + 12 * n, headed=True)
    r = dict(r or {})
    if r.get('planP99') is not None:
        sr = (r.get('starve') or 0) / max(1, r.get('ticks') or 1)
        r['starve_rate'] = round(sr, 4)
        pl = max(1, r.get('plans') or 1); why = r.get('certFailWhy') or {}
        r['cert_time_rate'] = round((why.get('time') or 0) / pl, 4); r['brake_rate'] = round((r.get('brakes') or 0) / pl, 4)
        w = r.get('planWarmP99') if r.get('planWarmP99') is not None else r['planP99']
        r['r9_pass'] = w < PLAN_LIMIT_MS and sr < STARVE_LIMIT and r['cert_time_rate'] < TIME_LIMIT
    print('viewer', json.dumps({k: v for k, v in r.items() if k != 'per'}, ensure_ascii=False), flush=True)
    return r


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('--model', default='ml/runs/r4/traj_policy_r4.onnx'); ap.add_argument('--out', default='ml/runs/r4.json')
    ap.add_argument('--n', type=int, default=150); ap.add_argument('--timeout', type=int, default=240)
    ap.add_argument('--viewer', action='store_true', help='只跑查看器 autotest（R9，bench 阶段用训练好的模型）')
    ap.add_argument('--viewer-n', type=int, default=30); ap.add_argument('--no-viewer-ref', action='store_true')
    a = ap.parse_args()
    if a.viewer: return main_viewer(a)
    out = {'rule': f'部署口径 batch=8 P99 < {LIMIT_MS:g} ms', 'model': a.model, 'pass': None}
    br = find_browser()
    if not br:
        out['note'] = '没找到 Edge / Chrome（可设 R4_BROWSER=浏览器路径）'; print(out['note'])
    else:
        out['browser'] = br
        srv = Srv()
        mu = '/' + os.path.relpath(os.path.abspath(a.model), ROOT).replace(os.sep, '/')
        runs = [one(srv, br, mu, 'wasm', 'auto', 1, a.n, a.timeout), one(srv, br, mu, 'wasm', '4', 1, a.n, a.timeout),
                one(srv, br, mu, 'wasm', 'auto', 0, a.n, a.timeout), one(srv, br, mu, 'webgpu', 'auto', 0, a.n, a.timeout)]
        out['runs'] = runs
        if not a.no_viewer_ref:   # 随机模型在查看器里的整次规划耗时：仅参考，不参与 R4 判定
            out['viewer_ref'] = viewer(srv, br, mu, a.viewer_n, True)
            out['viewer_ref_note'] = '随机权重：安全层路径不代表真模型，R9 以 bench 阶段 r9_browser.json 为准'
        srv.httpd.shutdown()
        dep, gpu = runs[0], runs[3]
        ok = lambda r: r.get('p99') is not None and r['p99'] < LIMIT_MS
        if dep.get('p99') is None and gpu.get('p99') is None:
            out['note'] = '页面没有给出数字（见 runs[].error）'
        elif ok(dep):
            out['pass'] = True; out['deploy'] = dep; out['deploy_ep'] = 'wasm'
        elif ok(gpu):
            out['pass'] = True; out['deploy'] = gpu; out['deploy_ep'] = 'webgpu'
            out['note'] = 'wasm 不达标而 WebGPU 达标：查看器部署用 &ep=webgpu'
        else:
            out['pass'] = False; out['deploy'] = dep if dep.get('p99') is not None else gpu
    os.makedirs(os.path.dirname(a.out) or '.', exist_ok=True)
    json.dump(out, open(a.out, 'w', encoding='utf-8'), indent=1, ensure_ascii=False)
    d = out.get('deploy') or {}
    print(f"R4：{ {True: '通过', False: '未通过', None: '未判定'}[out['pass']] }  部署口径 P99 {d.get('p99')} ms（{d.get('config')}）")


def main_viewer(a):
    meta = json.load(open(a.model.replace('.onnx', '.meta.json'), encoding='utf-8'))
    out = {'rule': f'R9：查看器稳态规划 planWarmP99 < {PLAN_LIMIT_MS:g} ms、运动中饥饿节拍 < {STARVE_LIMIT:.0%}、认证超时占比 < {TIME_LIMIT:.0%}（planColdP99 / brake_rate / tickGapMax 只报不判）', 'model': a.model, 'pass': None}
    br = find_browser()
    if not br: out['note'] = '没找到 Edge / Chrome（可设 R4_BROWSER=浏览器路径）'
    else:
        srv = Srv(); mu = '/' + os.path.relpath(os.path.abspath(a.model), ROOT).replace(os.sep, '/')
        r = viewer(srv, br, mu, a.viewer_n, bool(meta.get('random_init'))); srv.httpd.shutdown()
        out.update(browser=br, result=r)
        out['pass'] = r.get('r9_pass')
        if out['pass'] is None: out['note'] = r.get('error') or '页面没有给出整次规划耗时'
    os.makedirs(os.path.dirname(a.out) or '.', exist_ok=True)
    json.dump(out, open(a.out, 'w', encoding='utf-8'), indent=1, ensure_ascii=False)
    r = out.get('result') or {}
    print(f"R9：{ {True: '通过', False: '未通过', None: '未判定'}[out['pass']] }  稳态 planMs P99 {r.get('planWarmP99')} ms（冷启动 {r.get('planColdP99')}；推理 {r.get('inferP99')} / 安全层 {r.get('safetyP99')}），饥饿 {r.get('starve_rate')}，认证超时 {r.get('cert_time_rate')}，刹停 {r.get('brake_rate')}，tickGapMax {r.get('tickGapMax')}")


if __name__ == '__main__':
    main()