Electronic Load Molecule

#!/usr/bin/env python3
"""Local workcell agent/dashboard. Bind loopback; use SSH tunnel for remote use."""
import argparse,csv,io,json,threading,time,urllib.request,urllib.parse
from pathlib import Path
from http.server import ThreadingHTTPServer,BaseHTTPRequestHandler
from driver import Load
ROOT=Path(__file__).resolve().parent

def main():
 p=argparse.ArgumentParser();p.add_argument('--port',type=int,default=8864);p.add_argument('--resource');p.add_argument('--serial');p.add_argument('--simulate',action='store_true');p.add_argument('--control-panel',help='Existing local control-panel dashboard URL; read-only telemetry')
 p.add_argument('--origin',help='Exact allowed browser origin when a trusted proxy rewrites Host')
 a=p.parse_args()
 if a.simulate and a.resource:p.error('Simulation and hardware resource are mutually exclusive')
 load=Load(a.resource,a.serial,a.simulate);console=[];cp={'connected':False,'state':None};stop=threading.Event()
 def poll():
  while not stop.wait(.5):load.sample()
 def panel():
  while not stop.is_set():
   if a.control_panel:
    try:
     with urllib.request.urlopen(a.control_panel.rstrip('/')+'/api/state',timeout=2) as r:data=json.load(r)
     cp.update(connected=True,state=data,error=None,time=time.time())
    except Exception as e:cp.update(connected=False,state=None,error=str(e))
   stop.wait(2)
 threading.Thread(target=poll,daemon=True).start();threading.Thread(target=panel,daemon=True).start()
 class H(BaseHTTPRequestHandler):
  def log_message(self,*args):pass
  def send(self,data,status=200,typ='application/json'):
   b=json.dumps(data,allow_nan=False).encode() if typ=='application/json' else data
   self.send_response(status);self.send_header('Content-Type',typ);self.send_header('Content-Length',str(len(b)));self.send_header('Cache-Control','no-store');self.end_headers();self.wfile.write(b)
  def do_GET(self):
   path=urllib.parse.urlparse(self.path).path
   if path in ['/api/state','/state']:return self.send({**load.state(),'control_panel':cp,'history':load.samples[-240:]})
   if path=='/console':return self.send(console[-100:])
   if path=='/api/export.csv':
    rows=list(load.samples);s=io.StringIO()
    if rows:w=csv.DictWriter(s,fieldnames=list(rows[0]));w.writeheader();w.writerows(rows)
    return self.send(s.getvalue().encode(),typ='text/csv')
   if path=='/api/stream':
    self.send_response(200);self.send_header('Content-Type','text/event-stream');self.send_header('Cache-Control','no-cache');self.end_headers()
    try:
     for _ in range(600):
      self.wfile.write(('data: '+json.dumps({**load.state(),'control_panel':cp})+'\n\n').encode());self.wfile.flush();time.sleep(.5)
    except (BrokenPipeError,ConnectionResetError):pass
    return
   routes={'/':'ui/index.html','/index.html':'ui/index.html','/model.json':'3d-boards/preview.json','/provenance.json':'3d-boards/provenance.json','/three.js':'ui/three.js','/favicon.svg':'ui/favicon.svg','/model.step':'3d-boards/electronic-load-carrier.step','/model.glb':'3d-boards/electronic-load-carrier.glb'}
   if path not in routes:return self.send({'error':'Not found'},404)
   f=ROOT/routes[path];typ={'.html':'text/html','.json':'application/json','.svg':'image/svg+xml','.js':'text/javascript'}.get(f.suffix,'application/octet-stream')
   self.send(f.read_bytes(),typ=typ if typ!='application/json' else 'application/json; charset=utf-8')
  def do_POST(self):
   # JSON plus same-origin enforcement prevents a foreign website toggling localhost hardware.
   origin=self.headers.get('Origin')
   if origin and origin!=a.origin and urllib.parse.urlparse(origin).netloc!=self.headers.get('Host'):return self.send({'error':'Cross-origin control refused'},403)
   if self.headers.get('Content-Type','').split(';')[0]!='application/json':return self.send({'error':'JSON required'},415)
   try:
    n=int(self.headers.get('Content-Length','0'))
    if not 0<n<=8192:raise ValueError('Invalid request size')
    data=json.loads(self.rfile.read(n))
    if self.path=='/console':console.append(data);return self.send({'ok':True})
    if self.path!='/api/cmd':return self.send({'error':'Not found'},404)
    cmd=data.get('cmd')
    if cmd=='current':load.set_current(data.get('amps'))
    elif cmd=='on':load.on()
    elif cmd=='off':load.off()
    elif cmd=='heartbeat':load.heartbeat()
    else:raise ValueError('Unknown command')
    console.append({'time':time.time(),'cmd':cmd});self.send({'ok':True,'state':load.state()})
   except Exception as e:
    # Any failed hardware control attempts a confirmed OFF, never automatic retry of ON.
    if load.connected and not isinstance(e,ValueError):
     try:load.off()
     except Exception:load.error='Input state unknown; disconnect DUT power locally'
    self.send({'error':str(e)},400)
 server=ThreadingHTTPServer(('127.0.0.1',a.port),H)
 print(f'Electronic load: http://127.0.0.1:{a.port}/ | '+('SIMULATION' if a.simulate else 'HARDWARE / DISCONNECTED'),flush=True)
 try:server.serve_forever()
 finally:
  stop.set()
  if load.connected:
   try:load.off()
   except Exception:pass
  server.server_close()
if __name__=='__main__':main()