# SPDX-License-Identifier: Apache-2.0
"""End-to-end Ink Bench test in headless Chromium (Playwright for Python).
usage: python3 tests/e2e_bench.py <base url> <screenshot dir>
Drops a blurred, shifted and scaled copy of the published ink and pure noise through the UI,
checks alignment, scores and Python parity, annotation, export, share links, the leaderboard,
the cite panel, and a phone layout."""
import base64, json, os, sys, tempfile, zipfile
import numpy as np
from PIL import Image
from playwright.sync_api import sync_playwright
HERE = os.path.dirname(os.path.abspath(__file__)); ROOT = os.path.join(HERE, '..')
sys.path.insert(0, HERE); sys.path.insert(0, os.path.join(ROOT, 'tools'))
import synth, inkbench as ib
base, shots = sys.argv[1].rstrip('/') + '/', sys.argv[2]
os.makedirs(shots, exist_ok=True)
tmp = tempfile.mkdtemp()
CASES = {'shift_blur': ('w046', 23.0, -17.0, 1.0, 1.0, 2.0), 'scale_shift': ('w048', -12.5, 9.25, 1.12, 1.12, 1.0)}
for k, (seg, *p) in CASES.items():
    ib.save_png16(synth.make(seg, *p), os.path.join(tmp, k + '.png'))
Image.fromarray((synth.noise('w046') * 255).astype(np.uint8)).save(os.path.join(tmp, 'noise.png'))
refs = ib.Refs(root=ROOT)
R = {}; fails = []
def check(c, m):
    print(('ok   ' if c else 'FAIL ') + m)
    if not c: fails.append(m)
def scored(p):
    p.wait_for_function("window.unrollBench.state().scored && !window.unrollBench.state().busy", timeout=120000)
def comp(b):
    return ib.comparable(b)

with sync_playwright() as pw:
    br = pw.chromium.launch()
    ctx = br.new_context(viewport={'width': 1440, 'height': 900}, accept_downloads=True)
    p = ctx.new_page()
    errs = []
    p.on('pageerror', lambda e: errs.append(str(e)))
    p.goto(base, wait_until='load')
    p.wait_for_selector('#btnBench', timeout=60000)
    p.wait_for_selector('#loader', state='detached', timeout=120000)
    p.wait_for_timeout(500)
    p.click('#btnBench', timeout=180000)  # software WebGL in headless runs can keep the main thread busy at start
    p.wait_for_selector('#bench:not([hidden])')
    txt = p.inner_text('#bench')
    check('Nothing is uploaded' in txt, 'UI says nothing is uploaded')
    check('not ground truth' in txt, 'UI says agreement is not ground truth')
    # 1. blurred and shifted copy through the file input
    p.select_option('#bSeg', 'w046'); p.wait_for_timeout(500)
    p.fill('#bModel', 'synthetic: published ink, blur 2 px, shift 23,-17')
    p.set_input_files('#bPick', os.path.join(tmp, 'shift_blur.png'))
    scored(p)
    b = p.evaluate('window.unrollBench.bundle()')
    a, m = b['alignment'], b['metrics']
    R['shift_blur'] = {'truth': CASES['shift_blur'][1:], 'got': [a['tx'], a['ty'], a['sx'], a['sy'], m['blurSigmaPx']], 'pearson': m['pearson'], 'pearsonMatchedBlur': m['pearsonMatchedBlur'], 'iou': m['iou'], 'dice': m['dice'], 'ap': m['ap']}
    check(abs(a['tx'] - 23) <= 0.25 and abs(a['ty'] + 17) <= 0.25 and abs(a['sx'] - 1) < 0.003 and abs(a['sy'] - 1) < 0.003, f'shift recovered {a["tx"]},{a["ty"]} scale {a["sx"]},{a["sy"]}')
    check(m['blurSigmaPx'] == 2.0, f'blur recovered {m["blurSigmaPx"]} px')
    check(m['pearsonMatchedBlur'] > 0.99 and m['pearson'] > 0.85 and m['ap'] > 0.8, f'scores high r={m["pearson"]} rMatched={m["pearsonMatchedBlur"]} AP={m["ap"]}')
    py = ib.compare(os.path.join(tmp, 'shift_blur.png'), 'w046', model=b['prediction']['model'], refs=refs, name='shift_blur.png')
    check(comp(py) == comp(b), 'browser bundle identical to Python bundle (shift_blur)')
    p.screenshot(path=os.path.join(shots, 'bench_side.png'))
    p.click('[data-view="swipe"]'); p.wait_for_timeout(300)
    cv = p.locator('#bCanvases canvas').first.bounding_box()
    p.mouse.move(cv['x'] + cv['width'] * 0.5, cv['y'] + 50); p.mouse.down(); p.mouse.move(cv['x'] + cv['width'] * 0.4, cv['y'] + 60, steps=5); p.mouse.up()
    p.screenshot(path=os.path.join(shots, 'bench_swipe.png'))
    p.click('[data-view="diff"]'); p.wait_for_timeout(300)
    p.screenshot(path=os.path.join(shots, 'bench_diff.png'))
    # thresholds slider
    p.fill('#bTPred', '0.3'); p.dispatch_event('#bTPred', 'input'); p.wait_for_timeout(300)
    t = p.evaluate('window.unrollBench.metrics().thresholds')
    check(abs(t['pred'] - 0.3) < 1e-9, 'pred threshold slider updates metrics')
    b2 = p.evaluate('window.unrollBench.bundle()')
    py2 = ib.compare(os.path.join(tmp, 'shift_blur.png'), 'w046', model=b['prediction']['model'], refs=refs, name='shift_blur.png', t_pred=0.3)
    check(comp(py2) == comp(b2), 'identical to Python at pred threshold 0.3')
    p.fill('#bTPred', '0.5'); p.dispatch_event('#bTPred', 'input'); p.wait_for_timeout(200)
    # manual nudge
    p.click('[data-n="x+"]'); p.wait_for_timeout(600)
    na = p.evaluate('window.unrollBench.state()')
    check(abs(na['alignment']['tx'] - (a['tx'] + 1)) < 1e-6 and na['method'] == 'auto+manual', 'nudge moves 1 px and marks auto+manual')
    p.click('[data-n="x-"]'); p.wait_for_timeout(600)
    # annotation by mouse
    p.click('[data-view="side"]'); p.click('#bBox')
    cv = p.locator('#bCanvases canvas').first.bounding_box()
    p.mouse.move(cv['x'] + cv['width'] * 0.40, cv['y'] + cv['height'] * 0.40); p.mouse.down()
    p.mouse.move(cv['x'] + cv['width'] * 0.45, cv['y'] + cv['height'] * 0.47, steps=6); p.mouse.up()
    p.wait_for_timeout(200)
    p.select_option('#bAnn li select', 'Ε')
    p.fill('#bAnn li input[type=range]', '0.4'); p.dispatch_event('#bAnn li input[type=range]', 'input')
    ann = p.evaluate('window.unrollBench.annotations()')
    check(len(ann) == 1 and ann[0]['letter'] == 'Ε' and ann[0]['confidence'] == 0.4 and ann[0]['status'] == 'candidate for expert review', f'annotation drawn {ann}')
    check('candidate for expert review' in p.inner_text('#bAnn'), 'box list reads candidate for expert review')
    p.click('#bBox')
    p.screenshot(path=os.path.join(shots, 'bench_annotate.png'))
    # export
    with p.expect_download() as dl:
        p.click('#bExport')
    zp = os.path.join(tmp, dl.value.suggested_filename); dl.value.save_as(zp)
    z = zipfile.ZipFile(zp); names = sorted(n.split('/')[-1] for n in z.namelist())
    check(names == ['CITATION.cff', 'README.txt', 'bundle.json', 'figure.png', 'inkbench.bib', 'share-url.txt'], f'zip contents {names}')
    get = lambda suf: z.read([n for n in z.namelist() if n.endswith(suf)][0])
    zb = json.loads(get('bundle.json'))
    check(not ib.validate(zb), 'exported bundle passes inkbench validate')
    try:
        import jsonschema
        jsonschema.validate(zb, json.load(open(os.path.join(ROOT, 'docs/bench-bundle.schema.json')))); check(True, 'exported bundle passes the JSON Schema')
    except ImportError:
        pass
    pyz = ib.compare(os.path.join(tmp, 'shift_blur.png'), 'w046', model=zb['prediction']['model'], refs=refs, name='shift_blur.png', alignment=dict(zb['alignment'], method='auto+manual'), annotations=zb['annotations'])
    check(comp(pyz)['metrics'] == comp(zb)['metrics'] and comp(pyz)['annotations'] == comp(zb)['annotations'], 'Python reproduces the exported bundle metrics and annotations')
    fig = os.path.join(shots, 'bench_figure.png'); open(fig, 'wb').write(get('figure.png'))
    R['figure_px'] = Image.open(fig).size
    share_local = get('share-url.txt').decode().strip()
    # 2. noise
    p.set_input_files('#bPick', os.path.join(tmp, 'noise.png'))
    p.wait_for_timeout(300); scored(p)
    nb = p.evaluate('window.unrollBench.bundle()')
    nm = nb['metrics']
    R['noise'] = {k: nm[k] for k in ('pearson', 'pearsonMatchedBlur', 'iou', 'dice', 'ap')}; R['noise']['alignNcc'] = nb['alignment']['ncc']
    check(abs(nm['pearson']) < 0.1 and nm['ap'] < 0.2, f'noise scores low r={nm["pearson"]} AP={nm["ap"]}')
    check('low' in p.inner_text('#bAlignInfo'), 'noise alignment flagged as low confidence')
    pn = ib.compare(os.path.join(tmp, 'noise.png'), 'w046', model=nb['prediction']['model'], refs=refs, name='noise.png', annotations=nb['annotations'])
    check(comp(pn) == comp(nb), 'browser bundle identical to Python bundle (noise)')
    p.screenshot(path=os.path.join(shots, 'bench_noise.png'))
    # 3. scale case through the API with bytes
    b64 = base64.b64encode(open(os.path.join(tmp, 'scale_shift.png'), 'rb').read()).decode()
    p.evaluate('window.unrollBench.clearAnnotations()')
    sb = p.evaluate('async (b) => window.unrollBench.compare({ segment: "w048", bytes: b, name: "scale_shift.png", model: "synthetic scale" })', b64)
    sa = sb['alignment']
    R['scale_shift'] = {'truth': CASES['scale_shift'][1:], 'got': [sa['tx'], sa['ty'], sa['sx'], sa['sy'], sb['metrics']['blurSigmaPx']], 'pearson': sb['metrics']['pearson'], 'pearsonMatchedBlur': sb['metrics']['pearsonMatchedBlur']}
    check(abs(sa['tx'] + 12.5) <= 0.25 and abs(sa['ty'] - 9.25) <= 0.25 and abs(sa['sx'] / 1.12 - 1) < 0.003 and sb['metrics']['blurSigmaPx'] == 1.0, f'API scale case {R["scale_shift"]["got"]}')
    pys = ib.compare(os.path.join(tmp, 'scale_shift.png'), 'w048', model={'name': 'synthetic scale'}, refs=refs, name='scale_shift.png')
    check(comp(pys) == comp(sb), 'browser bundle identical to Python bundle (scale_shift via API)')
    # 4. public layer and share links
    gb = p.evaluate('async () => window.unrollBench.compare({ segment: "w049", layer: "gpu" })')
    url = p.evaluate('window.unrollBench.shareURL()')
    p2 = ctx.new_page(); p2.goto(url, wait_until='load'); scored(p2)
    rb = p2.evaluate('window.unrollBench.bundle()')
    check(rb['metrics'] == gb['metrics'] and rb['alignment']['tx'] == gb['alignment']['tx'], 'share link restores the public-layer comparison exactly')
    R['public_w049'] = {'pearson': gb['metrics']['pearson'], 'ap': gb['metrics']['ap'], 'shareURLchars': len(url)}
    p2.screenshot(path=os.path.join(shots, 'bench_share_restore.png')); p2.close()
    p3 = ctx.new_page(); p3.goto(share_local, wait_until='load'); p3.wait_for_selector('#bench:not([hidden])'); p3.wait_for_timeout(1500)
    check('Drop that same file' in p3.inner_text('#bStatus'), 'local-file share link asks for the same file')
    p3.set_input_files('#bPick', os.path.join(tmp, 'shift_blur.png')); scored(p3)
    lb = p3.evaluate('window.unrollBench.bundle()')
    check(lb['alignment']['method'] == 'shared' and lb['metrics'] == zb['metrics'] and len(lb['annotations']) == 1, 'local share link restores after dropping the same file')
    p3.close()
    # 5. leaderboard and cite
    p.click('[data-tab="board"]'); p.click('#bBoardEx'); p.wait_for_function("document.querySelectorAll('.b-table').length >= 37", timeout=180000); p.wait_for_timeout(800)
    rows = p.evaluate("[...document.querySelectorAll('.b-table')].map(t => [...t.querySelectorAll('tbody tr')].map(r => r.children[1].textContent))")
    check(len(rows) == 37 and all(r[0].startswith('scrollprize/ink_9um') and r[-1] == 'baseline: uniform noise' for r in rows), 'leaderboard ranks the 111 example bundles over 37 windings by AP, ink_9um first everywhere')
    rk = p.evaluate('async () => window.unrollBench.rank([window.unrollBench.bundle()], "ap")')
    check('w049' in rk, 'API rank works')
    p.screenshot(path=os.path.join(shots, 'bench_leaderboard.png'))
    p.click('[data-tab="cite"]'); p.wait_for_timeout(300)
    ct = p.inner_text('#bench')
    check('Apache License 2.0' in ct and 'CC BY-NC 4.0' in ct and 'Chase Hendrick' in ct, 'cite panel shows Apache-2.0, CC BY-NC 4.0 and attribution')
    p.screenshot(path=os.path.join(shots, 'bench_cite.png'))
    p.click('#bClose'); p.wait_for_timeout(200); p.click('#btnInfo', timeout=120000); p.wait_for_timeout(400)
    check('Why this matters for researchers' in p.inner_text('#about'), 'About panel has the researchers section')
    p.evaluate("document.querySelector('#about .bench-why').scrollIntoView()"); p.wait_for_timeout(500); p.screenshot(path=os.path.join(shots, 'bench_about.png'))
    R['api'] = p.evaluate("({ shim: typeof window.unrollBench, viaUnroll: !!(window.unroll && window.unroll.bench), hasUnroll: typeof window.unroll })")
    check(R['api']['shim'] == 'object' and (R['api']['hasUnroll'] != 'object' or R['api']['viaUnroll']), f'agent API {R["api"]}')
    check(not errs, f'no page errors {errs[:3]}')
    # 6. phone
    ph = br.new_context(**pw.devices['iPhone 13'])
    mp = ph.new_page(); merrs = []; mp.on('pageerror', lambda e: merrs.append(str(e)))
    mp.goto(base + '?bench', wait_until='load'); mp.wait_for_selector('#bench:not([hidden])', timeout=60000)
    mp.select_option('#bSeg', 'w046'); mp.wait_for_timeout(400)
    mp.set_input_files('#bPick', os.path.join(tmp, 'shift_blur.png')); scored(mp)
    mb = mp.evaluate('window.unrollBench.bundle()')
    check(comp(mb)['metrics'] == comp(py)['metrics'], 'phone gives the same metrics')
    cvb = mp.locator('#bCanvases canvas').first.bounding_box(); vw = mp.viewport_size['width']
    check(cvb['width'] <= vw and cvb['x'] >= 0, f'phone canvas fits the screen ({cvb["width"]:.0f} of {vw} px)')
    mp.tap('#bBox')
    mp.evaluate("""() => { const c = document.querySelector('#bCanvases canvas'), r = c.getBoundingClientRect();
      const ev = (t, x, y) => c.dispatchEvent(new PointerEvent(t, { bubbles: true, pointerId: 7, pointerType: 'touch', clientX: r.left + x * r.width, clientY: r.top + y * r.height }));
      ev('pointerdown', .3, .3); ev('pointermove', .33, .33); ev('pointermove', .38, .4); ev('pointerup', .38, .4); }""")
    mp.wait_for_timeout(200)
    check(len(mp.evaluate('window.unrollBench.annotations()')) == 1, 'phone touch draws a box')
    mp.screenshot(path=os.path.join(shots, 'bench_mobile.png'))
    mp.evaluate("document.querySelector('.b-metrics').scrollIntoView()"); mp.wait_for_timeout(300)
    mp.screenshot(path=os.path.join(shots, 'bench_mobile_metrics.png'))
    check(not merrs, f'no phone page errors {merrs[:3]}')
    br.close()
R['failures'] = fails
json.dump(R, open(os.path.join(shots, 'bench_results.json'), 'w'), indent=1)
print(json.dumps(R))
print('E2E', 'PASSED' if not fails else 'FAILED %d' % len(fails))
sys.exit(1 if fails else 0)