#!/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/native-browser-work');ROOT.mkdir(exist_ok=True);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=9334','--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:9334/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:8081/native_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(Path('/tmp/native-api-cases.json').read_text())
 parity=js('(async()=>{const cases='+json.dumps(cases,separators=(',',':'))+';for(const c of cases){const a=await NativeAltAI.simulate(NativeAltAI.model,Uint8Array.from(c.pixels));if(a.prediction!==c.prediction||!a.counts.every((v,i)=>v===c.hidden[i])||!a.output_counts.every((v,i)=>v===c.outputs[i]))throw Error("native parity "+c.index);}return{samples:cases.length,hidden_exact:true,output_exact:true};})()',True)
 def settle():
  return js("new Promise((resolve,reject)=>{let n=0;const t=setInterval(()=>{const answer=document.querySelector('#answer').textContent;if(answer!=='…'&&answer!=='—'){clearInterval(t);resolve({answer,result:FastMNIST.lastResult});}else if(document.querySelector('#error').textContent||++n>100){clearInterval(t);reject(Error(document.querySelector('#error').textContent||'timeout'));}},50);})",True)
 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
 js("document.querySelector('#gallery button').click()");settle();gallery=js("({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/'native-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})
 settle();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':[]})
 settle();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/'native-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}),500))",True)
 assert stable['blank'] and stable['answer']=='—',stable
 timing={'scope':'real server native KNP subprocess; not browser-only','milliseconds_gallery':gallery.get('milliseconds')}
 result={'ready':ready,'parity':parity,'blank':blank,'gallery2534':gallery,'mouse_drawn1':mouse,'touch_drawn1_mobile390':touch,'clear_after_debounce':stable,'inference_timing_headless':timing,'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)