import asyncio,json,sys,shutil
from pathlib import Path
import numpy as np
import torch
from PIL import Image
from huggingface_hub import hf_hub_download,HfApi
from coreai_mini import AmosMini

async def main():
 root=Path('bundles/amos-1-mini-CoreAI');source=Path('upstream/amos-1-mini').resolve()
 shutil.copytree(source/'serve',root/'serve',dirs_exist_ok=True)
 for f in ['config.json','LICENSE']:shutil.copy2(source/f,root/f)
 from convert_mini import PINNED_DEPENDENCIES
 cfg=json.loads((source/'config.json').read_text());deps=[]
 for repo in [cfg['repo'],cfg['vision']]:
  revision=PINNED_DEPENDENCIES[repo];deps.append({'repo':repo,'revision':revision})
  if repo==cfg['repo']:shutil.copy2(hf_hub_download(repo,'tokenizer.json',revision=revision),root/'tokenizer.json')
 (root/'DEPENDENCIES.json').write_text(json.dumps(deps,indent=2))
 sys.path.insert(0,str(source/'serve'))
 from amos.model import load_run,pack_request,decode
 from amos.packing import collate
 from amos.vision import preprocess
 from convert_mini import pin_dependencies
 pin_dependencies()
 torch.set_num_threads(4);reference,codec=load_run(source)
 ai=await AmosMini.load(root)
 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)]}],[]),
 ]
 pixels=np.zeros((224,224,3),np.uint8);pixels[:,:112]=[255,0,0];pixels[:,112:]=[0,0,255];im=Image.fromarray(pixels)
 cases += [('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=[]
 for name,state,qs,ims in cases:
  refqs=[{**q,'options':[preprocess(o) if isinstance(o,Image.Image) else o for o in q['options']]} if 'options' in q else q for q in qs]
  with torch.no_grad():
   ps=pack_request(reference,codec,state,refqs,[preprocess(x) for x in ims]);expected=reference(collate(ps,codec.pad));ed=decode(expected,ps)[0]
  actual,aps=await ai.raw(state,qs,ims);ad=decode(actual,aps)[0]
  errors={}
  for key in ['opt','span','region']:
   if key not in expected:continue
   x=expected[key];y=actual[key];mask=torch.isfinite(x)
   if not torch.equal(mask,torch.isfinite(y)):raise AssertionError((name,key,x.shape,y.shape,x,y))
   errors[key]=float((x[mask]-y[mask]).abs().max()) if mask.any() else 0
   if errors[key]>0.005:raise AssertionError((name,key,errors))
  def decisions(rs):
   return [{k:v for k,v in r.items() if k in ['type','text','start','end','cell','choice','index']} for r in rs]
  if decisions(ed)!=decisions(ad):raise AssertionError((name,ed,ad))
  rec={'case':name,'tokens':len(ps[0]),'max_logit_error':errors,'upstream':ed,'coreai':ad};records.append(rec);print(json.dumps(rec),flush=True)
 (root/'validation.json').write_text(json.dumps(records,indent=2));print('PASS',flush=True)
if __name__=='__main__':asyncio.run(main())