#!/usr/bin/env python3
"""Bounded Chromium CDP parity and real mouse/touch smoke test."""
import base64,json,os,shutil,socket,subprocess,tempfile,time,urllib.request
from pathlib import Path
import websocket
ROOT=Path('/tmp/mnist-fast-work');profile=tempfile.mkdtemp(prefix='mnist-cdp-');errors=[];counter=0
log=(ROOT/'browser-check.log').open('w');browser=subprocess.Popen(['chromium','--headless','--no-sandbox','--disable-gpu','--remote-debugging-port=9333','--remote-allow-origins=*','--user-data-dir='+profile,'about:blank'],stdout=log,stderr=log)
ws=None
try:
 for _ in range(50):
  try:
   pages=json.load(urllib.request.urlopen('http://127.0.0.1:9333/json',timeout=1));break
  except Exception:time.sleep(.1)
 else:raise RuntimeError('Chromium CDP unavailable')
 ws=websocket.create_connection(next(p['webSocketDebuggerUrl'] for p in pages if p['type']=='page'),timeout=20)
 def call(method,params=None):
  global counter
  counter+=1;want=counter;ws.send(json.dumps({'id':want,'method':method,'params':params or {}}))
  while True:
   message=json.loads(ws.recv())
   if message.get('method')=='Runtime.exceptionThrown':errors.append(message['params'])
   if message.get('id')==want:
    if 'error' in message:raise RuntimeError(message['error'])
    return message.get('result',{})
 def js(expression,await_promise=False):
  result=call('Runtime.evaluate',{'expression':expression,'returnByValue':True,'awaitPromise':await_promise})
  if 'exceptionDetails' in result:raise RuntimeError(result['exceptionDetails'])
  return result['result'].get('value')
 call('Page.enable');call('Runtime.enable');call('Page.navigate',{'url':'http://10.9.8.230:8080/fast_viz.html'})
 ready=js("new Promise((resolve,reject)=>{let n=0;const t=setInterval(()=>{if(window.FastMNIST){clearInterval(t);FastMNIST.ready.then(m=>resolve({version:m.version,steps:m.steps})).catch(reject);}else if(++n>100){clearInterval(t);reject(Error('page not ready'));}},50);})",True)
 cases=json.loads((ROOT/'parity-cases.json').read_text())['cases']
 parity=js('''(()=>{const cases='''+json.dumps(cases,separators=(',',':'))+''';let normalized=0,counts=0,scores=0,predictions=0;for(const item of cases){const decode=s=>Uint8Array.from(atob(s),x=>x.charCodeAt(0));const pixels=FastMNIST.normalizeGray(decode(item.raw));if(!pixels.every((v,i)=>v===decode(item.normalized)[i]))throw Error('normalization '+item.test_index);normalized++;const out=FastMNIST.simulate(FastMNIST.model,pixels);if(!out.counts.every((v,i)=>v===item.counts[i]))throw Error('IF counts '+item.test_index);counts++;if(!out.scores.every((v,i)=>v===item.scores[i]))throw Error('scores '+item.test_index);scores++;if(out.prediction!==item.prediction)throw Error('prediction '+item.test_index);predictions++;}return{normalized,counts,scores,predictions};})()''')
 blank=js("(()=>{document.querySelector('#clear').click();return {answer:document.querySelector('#answer').textContent,blank:FastMNIST.lastResult.blank,scores:document.querySelector('#scores').children.length};})()")
 assert blank=={'answer':'—','blank':True,'scores':0},blank
 gallery=js("(()=>{document.querySelector('#gallery button').click();return{answer:document.querySelector('#answer').textContent,prediction:FastMNIST.lastResult.prediction,sample:document.querySelector('#sample').textContent};})()")
 assert gallery['prediction']==3,gallery
 shot=call('Page.captureScreenshot',{'format':'png','captureBeyondViewport':True});(ROOT/'fast-viz-desktop.png').write_bytes(base64.b64decode(shot['data']))
 # Mouse: draw a simple1 with genuine browser inputevents.
 js("document.querySelector('#clear').click();document.querySelector('#draw').scrollIntoView({block:'center'});")
 rect=js("(()=>{let r=document.querySelector('#draw').getBoundingClientRect();return{x:r.x,y:r.y,w:r.width,h:r.height}})()")
 x=rect['x']+.52*rect['w'];y0=rect['y']+.18*rect['h'];y1=rect['y']+.82*rect['h']
 call('Input.dispatchMouseEvent',{'type':'mousePressed','x':x,'y':y0,'button':'left','clickCount':1})
 for i in range(1,16):call('Input.dispatchMouseEvent',{'type':'mouseMoved','x':x,'y':y0+(y1-y0)*i/15,'button':'left','buttons':1})
 call('Input.dispatchMouseEvent',{'type':'mouseReleased','x':x,'y':y1,'button':'left','clickCount':1})
 mouse=js("({prediction:FastMNIST.lastResult.prediction,blank:FastMNIST.lastResult.blank,ink:FastMNIST.lastNormalized.reduce((a,b)=>a+b,0)})")
 assert mouse['prediction']==1 and not mouse['blank'],mouse
 call('Emulation.setDeviceMetricsOverride',{'width':390,'height':844,'deviceScaleFactor':1,'mobile':True})
 call('Emulation.setTouchEmulationEnabled',{'enabled':True,'maxTouchPoints':1})
 js("document.querySelector('#clear').click();document.querySelector('#draw').scrollIntoView({block:'center'});")
 rect=js("(()=>{let r=document.querySelector('#draw').getBoundingClientRect();return{x:r.x,y:r.y,w:r.width,h:r.height}})()")
 x=rect['x']+.5*rect['w'];y0=rect['y']+.18*rect['h'];y1=rect['y']+.82*rect['h']
 call('Input.dispatchTouchEvent',{'type':'touchStart','touchPoints':[{'x':x,'y':y0,'id':1}]})
 for i in range(1,16):call('Input.dispatchTouchEvent',{'type':'touchMove','touchPoints':[{'x':x,'y':y0+(y1-y0)*i/15,'id':1}]})
 call('Input.dispatchTouchEvent',{'type':'touchEnd','touchPoints':[]})
 touch=js("({prediction:FastMNIST.lastResult.prediction,blank:FastMNIST.lastResult.blank,ink:FastMNIST.lastNormalized.reduce((a,b)=>a+b,0),overflow:document.documentElement.scrollWidth>innerWidth})")
 assert touch['prediction']==1 and not touch['blank'] and not touch['overflow'],touch
 shot=call('Page.captureScreenshot',{'format':'png','captureBeyondViewport':True});(ROOT/'fast-viz-mobile.png').write_bytes(base64.b64decode(shot['data']))
 # Clearing immediately afterpointer inputmustsurvive an olddebounce timer.
 js("document.querySelector('#clear').click()")
 stable=js("new Promise(r=>setTimeout(()=>r({blank:FastMNIST.lastResult.blank,answer:document.querySelector('#answer').textContent}),150))",True)
 assert stable['blank'] and stable['answer']=='—',stable
 timing=js("(()=>{let x=Uint8Array.from(atob('"+cases[0]["normalized"]+"'),c=>c.charCodeAt(0));let times=[];for(let i=0;i<120;i++){const r=FastMNIST.simulate(FastMNIST.model,x);if(i>=20)times.push(r.milliseconds);}times.sort((a,b)=>a-b);return{samples:100,median_ms:times[50],p95_ms:times[95]};})()")
 result={'ready':ready,'parity':parity,'blank':blank,'gallery2534':gallery,'mouse_drawn1':mouse,'touch_drawn1_mobile390':touch,'clear_after_debounce':stable,'inference_timing_headless':{**timing,'sample':'officialtest2534normalized','scope':'networkonly; excludescanvasnormalization/DOM','hardware':'localheadlessChromiumnotuserdevice'},'javascript_errors':errors}
 assert not errors,errors
 (ROOT/'browser-check.json').write_text(json.dumps(result,indent=2)+'\n');print(json.dumps(result))
finally:
 if ws:ws.close()
 browser.terminate()
 try:browser.wait(timeout=5)
 except subprocess.TimeoutExpired:browser.kill();browser.wait()
 log.close();shutil.rmtree(profile,ignore_errors=True)