import argparse,asyncio,json,sys,shutil
from pathlib import Path
import numpy as np
import torch
from PIL import Image
from coreai_decoder import AmosDecoder

async def main(name,mode='both'):
 source=(Path('upstream')/name).resolve();root=Path('bundles')/(name+'-CoreAI')
 shutil.copytree(source/'serve',root/'serve',dirs_exist_ok=True)
 (root/'model').mkdir(exist_ok=True)
 for f in (source/'model').iterdir():
  if f.suffix not in ['.safetensors','.pt']:shutil.copy2(f,root/'model'/f.name)
 for f in ['config.json','calibration.json','LICENSE']:shutil.copy2(source/f,root/f)
 sys.path.insert(0,str(source/'serve'))
 import amos.decoder as decoder
 from transformers import Qwen2VLImageProcessorPil
 decoder.IMAGE_PIXELS=224*224
 precision=json.loads((root/'manifest.json').read_text())['precision']
 torch.set_num_threads(4)
 if mode!='coreai':
  model,codec=decoder.load_model(str(source/'model'),'cpu',dtype=getattr(torch,precision));model.eval()
  codec.image_processor=Qwen2VLImageProcessorPil.from_pretrained(source/'model',size={'shortest_edge':224*224,'longest_edge':224*224})
  heads=decoder.load_heads(source/'model',model,codec);heads.span=decoder.span_decoding(source)
 if mode!='reference':ai=await AmosDecoder.load(root)
 pixels=np.zeros((224,224,3),np.uint8);pixels[:,:112]=[255,0,0];pixels[:,112:]=[0,0,255];im=Image.fromarray(pixels)
 cases=[
 ('refund','Invoice INV-2231 for £306.00 was charged twice. Please refund one.',[{'type':'noul','text':'Is a refund requested?'},{'type':'choice','text':'Which team?','options':['billing','technical','sales']}],[]),
 ('span','The invoice total is £306.00 and the invoice number is INV-2231.',[{'type':'span','text':'Quote the invoice number.'}],[]),
 ('score','The service is down for every customer.',[{'type':'score','text':'How severe is the outage?','options':['none','minor','major','complete']}],[]),
 ('multi','I need a refund and a password reset.',[{'type':'multi','text':'Which topics apply?','options':['refund','password','delivery']}],[]),
 ('negative','The customer explicitly says no refund is needed.',[{'type':'noul','text':'Is a refund requested?'}],[]),
 ('unicode','José paid €123.45 for café supplies.',[{'type':'span','text':'Quote the amount paid.'},{'type':'choice','text':'Who paid?','options':['José','Maria','Bob']}],[]),
 ('options','This is about category 17.',[{'type':'choice','text':'Which category?','options':[f'category {i}' for i in range(24)]}],[]),
 ('vision','Two colored halves.',[{'type':'region','text':'Point to the red half.'},{'type':'choice','text':'Which colors?','options':['red and blue','green and yellow']}],[im]),
 ('image-option','Find the red-blue image.',[{'type':'choice','text':'Which image matches?','options':[im,Image.new('RGB',(224,224),'green')]}],[]),
 ]
 records=[];refs=json.loads((root/'reference.json').read_text()) if mode=='coreai' else []
 for case,state,qs,ims in cases:
  if mode=='coreai':
   ref=next(r for r in refs if r['case']==case);expected,tokens=ref['upstream'],ref['tokens']
  else:
   with torch.no_grad():expected,tokens=decoder.results(model,codec,state,qs,ims,heads,512,temperature=decoder.calibration(source))
  if mode=='reference':
   records.append({'case':case,'tokens':tokens,'reference_precision':precision,'upstream':expected});print(case,tokens,flush=True);continue
  actual=await ai.predict(state,qs,ims)
  error=0;match=True
  for x,y in zip(expected,actual):
   for key in ['p_yes','probs','p','box']:
    if key in x:error=max(error,float(np.abs(np.asarray(x[key])-np.asarray(y[key])).max()))
   for key in ['choice','start','end','text','cell']:
    if key in x and x[key]!=y[key]:match=False
  rec={'case':case,'tokens':tokens,'reference_precision':precision,'max_probability_or_box_error':error,'decisions_match':match,'upstream':expected,'coreai':actual};records.append(rec);print(json.dumps(rec),flush=True)
  if not match or error>0.02:raise AssertionError(rec)
 (root/('reference.json' if mode=='reference' else 'validation.json')).write_text(json.dumps(records,indent=2));print('PASS',flush=True)
if __name__=='__main__':
 p=argparse.ArgumentParser();p.add_argument('name');p.add_argument('--mode',choices=['both','reference','coreai'],default='both');a=p.parse_args();asyncio.run(main(a.name,a.mode))