"""方案 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` 天生不完整）。
范围说明：这里测的是纯推理；贴柱处安全层耗时（R9）由 bench（tests/policy_bench.cjs 的 safetyMsP99）给。"""
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


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 one(srv, browser, model_url, ep, threads, proxy, n, timeout_s):
    srv.result = None; srv.ev.clear()
    q = urllib.parse.urlencode({'model': model_url, 'ep': ep, 'threads': threads, 'proxy': proxy, 'n': n,
                                **({'ortLocal': '1'} if os.environ.get('ORT_LOCAL') == '1' else {})})
    url = f'http://127.0.0.1:{srv.port}/webapp/r4_probe.html?{q}'
    prof = tempfile.mkdtemp(prefix='r4prof_')
    args = [browser, '--headless=new', f'--user-data-dir={prof}', '--no-first-run', '--no-default-browser-check',
            '--disable-extensions', '--enable-unsafe-webgpu', 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)
    r = srv.result if got else {'error': f'{timeout_s} s 内页面没有回报'}
    r = dict(r or {}, config=f'{ep}·threads={threads}·proxy={proxy}')
    print(json.dumps(r, 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)
    a = ap.parse_args()
    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)]
        srv.httpd.shutdown()
        out['runs'] = runs
        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')}）")


if __name__ == '__main__':
    main()