"""Text-only Core AI Decision adapter. Adapted from Max Farrell, Apache-2.0.
Changes by hugging-mac: bounded text-only input, explicit busy rejection, lazy-free
one-request lifetime, original question normalization and FP32 softmax.
"""
import json
from pathlib import Path
import numpy as np
from coreai.runtime import AIModel,NDArray,SpecializationOptions
from transformers import AutoTokenizer
from d3_format import render
from d3_runtime import normalize_question,product_answer

def pack(encoded,manifest):
    ids=np.asarray(encoded['input_ids'],dtype=np.int32)
    n=ids.shape[1];length=manifest['sequence_length']
    if ids.shape[0]!=1 or n<1:raise ValueError('Exactly one nonempty prompt is required')
    if n>length:raise ValueError(f'Prompt needs {n} tokens; capacity is {length}; input was not truncated')
    pad=length-n
    # Text RoPE has translation-invariant positions; real tokens start at zero.
    positions=np.zeros((4,1,length),dtype=np.int32)
    positions[:,:,pad:]=np.arange(n,dtype=np.int32)
    return dict(input_ids=np.pad(ids,((0,0),(pad,0))),attention_mask=np.pad(np.ones((1,n),dtype=np.int32),((0,0),(pad,0))),position_ids=positions)

class D3CoreAI:
    @classmethod
    async def load(cls,bundle):
        self=cls();self.root=Path(bundle)
        self.manifest=json.loads((self.root/'manifest.json').read_text())
        self.config=json.loads((self.root/'decision_config.json').read_text())
        self.tokenizer=AutoTokenizer.from_pretrained(self.root,local_files_only=True)
        self.model=await AIModel.load(self.root/'decision.aimodel',SpecializationOptions.default())
        self.fn=self.model.load_function('main');self._busy=False
        return self

    async def logits(self,state,question,images=None,videos=None):
        if images or videos:raise ValueError('This artifact supports text only')
        if self._busy:raise RuntimeError('Decision inference is busy; no implicit waiting queue')
        normalized=normalize_question(question)
        text=render(self.tokenizer,state,normalized.rendered,self.config['codes'])
        encoded=self.tokenizer([text],add_special_tokens=False,return_tensors='np')
        arrays=pack(encoded,self.manifest)
        self._busy=True
        try:
            output=await self.fn({k:NDArray(v) for k,v in arrays.items()})
            logits=output[self.manifest['output']].numpy()[0,:len(normalized.keys)]
            if not np.isfinite(logits).all():raise FloatingPointError('Nonfinite Core AI logits')
            return logits,normalized.keys,encoded
        finally:self._busy=False

    async def system_one(self,*,state,questions,images=None,videos=None):
        if not isinstance(questions,dict) or not 1<=len(questions)<=16:raise ValueError('Provide 1 to 16 questions')
        answers={}
        for key,q in questions.items():
            logits,keys,_=await self.logits(state,q,images,videos)
            z=logits.astype(np.float32)/self.config['temperature'];probs=np.exp(z-z.max());probs/=probs.sum()
            if not np.isfinite(probs).all():raise FloatingPointError('Nonfinite probabilities')
            answers[key]=product_answer(normalize_question(q),probs.tolist())
        return {'model':self.manifest['source'],'answers':answers}