"""V3 experimental intent inference wrapper. No networking or automatic model downloads."""
import argparse,hashlib,json,time
from pathlib import Path
import torch
from safetensors.torch import load_file
from transformers import AutoConfig,AutoModel,AutoTokenizer
from laya_common import DecisionModel,build_sequence,collate_items,render_options,serialize_state,QTYPES

def file_hash(path):
 h=hashlib.sha256()
 with path.open('rb') as f:
  for chunk in iter(lambda:f.read(8*1024*1024),b''):h.update(chunk)
 return h.hexdigest()

class Agent:
 def __init__(self,root=None,device='cpu'):
  self.root=Path(root) if root else Path(__file__).resolve().parent
  self.manifest=json.loads((self.root/'manifest.json').read_text())
  if self.manifest['status']!='verified':raise ValueError('Model artifact download and verification not complete')
  for entry in self.manifest['files']:
   if file_hash(self.root/entry['path'])!=entry['sha256']:raise ValueError('Artifact checksum mismatch: '+entry['path'])
  if device not in ['cpu','mps','cuda']:raise ValueError('Unsupported device')
  if device=='mps' and not torch.backends.mps.is_available():raise RuntimeError('MPS unavailable')
  if device=='cuda' and not torch.cuda.is_available():raise RuntimeError('CUDA unavailable')
  self.device=torch.device(device);self.tok=AutoTokenizer.from_pretrained(self.root/'tokenizer',local_files_only=True)
  cfg=AutoConfig.from_pretrained(self.root/'encoder',local_files_only=True);cfg.reference_compile=False
  self.model=DecisionModel(AutoModel.from_config(cfg,attn_implementation='sdpa'),2,2)
  self.model.load_state_dict(load_file(str(self.root/'model.safetensors')),strict=True)
  self.model.to(self.device).eval()

 def sync(self):
  if self.device.type=='mps':torch.mps.synchronize()
  elif self.device.type=='cuda':torch.cuda.synchronize()

 def predict(self,state,questions):
  if not questions:raise ValueError('At least one question required')
  if len(questions)>16:raise ValueError('At most 16 questions per call')
  self.sync();start=time.perf_counter();items=[];details=[]
  for qid,q in questions.items():
   kind=q.get('type');ins=q.get('instructions')
   if kind not in ['choice','noul','score'] or not isinstance(ins,str) or not ins.strip():raise ValueError('Invalid question type/instructions')
   crit=q.get('criteria')
   if kind!='noul' and (not isinstance(crit,(dict,list)) or len(crit)<2):raise ValueError('At least two choice/score criteria required')
   query={'t':kind,'ins':ins,'crit':crit}
   opts=render_options(query)
   ids,markers=build_sequence(self.tok,state,query,1024,512)
   empty,_=build_sequence(self.tok,'',query,1024,512)
   if len(markers)!=len(opts):raise ValueError('Question exceeds candidate budget')
   if ids.index(self.tok.sep_token_id)-1!=len(self.tok(f'{kind} question: {ins}',add_special_tokens=False)['input_ids']):raise ValueError('Instruction exceeds token budget')
   boundary=next(i for i in range(markers[-1]+1,len(ids)) if ids[i]==self.tok.sep_token_id)
   tokens=[tuple(ids[m+1:(markers[i+1] if i+1<len(markers) else boundary)]) for i,m in enumerate(markers)]
   if len(set(tokens))!=len(opts):raise ValueError('Candidate text collision after truncation')
   source_len=len(self.tok(serialize_state(state).replace(self.tok.mask_token,' '),add_special_tokens=False)['input_ids'])
   details.append({'question':qid,'state_truncated':len(ids)-len(empty)<source_len,'input_tokens':len(ids)})
   items.append({'ids':ids,'markers':markers,'qtype':QTYPES[kind]})
  batch=collate_items([items],self.tok.pad_token_id);x={k:batch[k].to(self.device) for k in ['input_ids','attention_mask','marker_pos','marker_mask','qtype']}
  with torch.inference_mode():logits,_=self.model(**x)
  answers={}
  for i,(qid,q) in enumerate(questions.items()):
   p=logits[i,:len(items[i]['markers'])].float().softmax(-1).cpu()
   if not torch.isfinite(p).all():raise RuntimeError('Nonfinite prediction')
   values=p.tolist();kind=q['type']
   if kind=='noul':answers[qid]={'noul':values[1]}
   else:
    labels=list(q['criteria']);answer={'probabilities':dict(zip(labels,values))}
    if kind=='choice':answer['choice']=labels[int(p.argmax())]
    else:answer['score']=sum(i*v for i,v in enumerate(values))
    answers[qid]=answer
  self.sync()
  return {'answers':answers,'elapsed_ms':(time.perf_counter()-start)*1000,'device':str(self.device),'dtype':'float32','diagnostics':details,'calibration':'uncalibrated; no trained abstention/action guarantee'}

if __name__=='__main__':
 p=argparse.ArgumentParser();p.add_argument('--input',required=True);p.add_argument('--device',choices=['cpu','mps','cuda'],default='cpu');args=p.parse_args();payload=json.loads(Path(args.input).read_text());print(json.dumps(Agent(device=args.device).predict(payload['state'],payload['questions']),ensure_ascii=False,indent=2))