"""CIFAR pets demo. All learned inference is KNP; host only image processing, fixed pooling/packet routing."""
from http.server import ThreadingHTTPServer,BaseHTTPRequestHandler
from pathlib import Path
from PIL import Image,ImageOps
import base64,io,json,threading,tempfile,subprocess,time,csv,warnings,hashlib
import numpy as np
from worker import DetectorWorker
WORKER=DetectorWorker()
S=Path(__file__).resolve().parent;R=Path.home()/'run/pets-knp';B=Path.home()/'KNP/build/bin';LOCK=threading.Lock();Image.MAX_IMAGE_PIXELS=16_000_000
DATA=np.load(R/'data/test.npz'); samples=[]
for c in (0,1):samples.extend(np.flatnonzero(DATA['labels']==c)[:12].tolist())
samples=sorted(samples)
def encoded(im):
 b=io.BytesIO();im.save(b,format='PNG');return 'data:image/png;base64,'+base64.b64encode(b.getvalue()).decode()
MODELINFO=json.loads((R/'final-summary.json').read_text())
def metadata():return MODELINFO
def verify_models():
 for name,sha in MODELINFO['files_sha256'].items():
  if hashlib.sha256((R/name).read_bytes()).hexdigest()!=sha:raise RuntimeError('Frozen model was changed: '+name)
class Handler(BaseHTTPRequestHandler):
 def send(self,code,value,ctype='application/json; charset=utf-8'):
  raw=json.dumps(value,ensure_ascii=False).encode() if not isinstance(value,bytes) else value
  self.send_response(code);self.send_header('Content-Type',ctype);self.send_header('Content-Length',str(len(raw)));self.send_header('Cache-Control','no-store');self.send_header('X-Content-Type-Options','nosniff');self.end_headers();self.wfile.write(raw)
 def do_GET(self):
  if self.path in ('/','/pets.html'):return self.send(200,(S/'pets.html').read_bytes(),'text/html; charset=utf-8')
  if self.path=='/api/info':return self.send(200,{k:MODELINFO[k] for k in ['version','test_accuracy','test_total','model_sha256','pipeline_sha256','classes']})
  if self.path=='/api/samples':return self.send(200,[{'id':i,'label':int(DATA['labels'][i]),'image':encoded(Image.fromarray(DATA['rgb'][i].transpose(1,2,0)))} for i in samples])
  self.send(404,{'error':'Не найдено'})
 def do_POST(self):
  if self.path!='/api/recognize':return self.send(404,{'error':'Не найдено'})
  if not LOCK.acquire(blocking=False):return self.send(429,{'error':'GPU занят предыдущей картинкой. Попробуйте через несколько секунд.'})
  try:self.recognize()
  finally:LOCK.release()
 def recognize(self):
  try:
   size=int(self.headers.get('Content-Length','0'))
   if not 0<size<=12_000_000:return self.send(413,{'error':'Слишком большой файл. Максимум 8 МБ.'})
   body=json.loads(self.rfile.read(size));uri=body.get('image')
   if not isinstance(uri,str) or not uri.startswith('data:image/') or ';base64,' not in uri:raise ValueError('Ожидается изображение')
   raw=base64.b64decode(uri.split(',',1)[1],validate=True)
   if len(raw)>8_000_000:return self.send(413,{'error':'Максимум 8 МБ'})
   with warnings.catch_warnings():
    warnings.simplefilter('error',Image.DecompressionBombWarning)
    image=Image.open(io.BytesIO(raw))
    if image.format not in ('PNG','JPEG','WEBP','BMP'):raise ValueError('Используйте PNG, JPEG, WebP или BMP')
    if image.width*image.height>16_000_000:raise ValueError('Изображение больше 16 мегапикселей')
    image=ImageOps.exif_transpose(image).convert('RGB').resize((32,32),Image.Resampling.BILINEAR)
   px=np.asarray(image).transpose(2,0,1).copy()
  except Exception as e:return self.send(400,{'error':'Не удалось прочитать изображение: '+str(e)[:120]})
  start=time.monotonic()
  try:
   verify_models()
   with tempfile.TemporaryDirectory(prefix='request-',dir=R) as d:
    d=Path(d);px.tofile(d/'image.bin')
    WORKER.infer(d/'image.bin',d/'det')
    counts=np.fromfile(d/'det_counts.bin',np.int32)
    # Fixed root-L2 contrast normalization and count routing, never learned arithmetic.
    q=np.sqrt(counts.astype(np.float64));q*=2048/max(float(np.sqrt((q*q).sum())),1)
    np.clip(np.rint(q),0,640).astype(np.int32).tofile(d/'encoded.bin')
    with (d/'hidden.log').open('w') as log:subprocess.run([str(R/'pets_hidden'),str(d/'encoded.bin'),str(R/'head-v2'),'cuda',str(d/'hidden.bin')],stdout=log,stderr=subprocess.STDOUT,check=True,timeout=15)
    with (d/'readout.log').open('w') as log:subprocess.run([str(R/'pets_readout'),str(d/'hidden.bin'),str(R/'head-v2'),'cuda',str(d/'result.csv')],stdout=log,stderr=subprocess.STDOUT,check=True,timeout=15)
    row=next(csv.DictReader((d/'result.csv').open()));counts=np.fromfile(d/'det_counts.bin',np.int32);heat=np.fromfile(d/'det_heat.bin',np.int32).tolist()
    result={'prediction':int(row['prediction']),'outputs':[int(row['cat']),int(row['dog'])],'tie':int(row['cat'])==int(row['dog']),'input':encoded(image),'heat':heat,'heatSide':14,'detectorSpikes':int(counts.sum()),'milliseconds':round((time.monotonic()-start)*1000),'backend':'KNP CUDA / AltAILIF','model':metadata()['model_sha256']}
   self.send(200,result)
  except subprocess.TimeoutExpired:self.send(504,{'error':'Превышено время расчёта. Повторите запрос.'})
  except Exception as e:
   print('INFERENCE ERROR',repr(e),flush=True);self.send(500,{'error':'Ошибка KNP. Подробности в журнале сервера.'})
class Server(ThreadingHTTPServer):
 def __init__(self,*args,**kwargs):self.admission=threading.BoundedSemaphore(4);super().__init__(*args,**kwargs)
 def process_request(self,request,address):
  request.settimeout(15)
  if not self.admission.acquire(blocking=False):
   try:request.sendall(b'HTTP/1.0 503 Service Unavailable\r\nContent-Length: 0\r\nConnection: close\r\n\r\n')
   finally:self.shutdown_request(request)
   return
  try:super().process_request(request,address)
  except Exception:self.admission.release();raise
 def process_request_thread(self,request,address):
  try:super().process_request_thread(request,address)
  finally:self.admission.release()
if __name__=='__main__':
 verify_models()
 # Warm graph before announcing readiness. No request/model is fabricated or cached.
 with tempfile.TemporaryDirectory(prefix='warmup-',dir=R) as d:
  d=Path(d);np.zeros((3,32,32),np.uint8).tofile(d/'black.bin');WORKER.infer(d/'black.bin',d/'black')
 print('READY pets server :8082',flush=True)
 Server(('0.0.0.0',8082),Handler).serve_forever()