"""The confidence layer text reasoning on immutable answer-contract base-prefix states."""
import copy,hashlib,json
from pathlib import Path
from solomon.engine_numerics import CudaEngine

FORCE='\n</think>\n\nGive only the answer letter: '

def validate_request(state,n_letters,max_tokens):
    if not isinstance(max_tokens,int) or isinstance(max_tokens,bool) or not 1<=max_tokens<=3072:raise ValueError('reasoning cap must be 1..3072')
    if not isinstance(n_letters,int) or isinstance(n_letters,bool) or not 2<=n_letters<=10:raise ValueError('invalid letter count')
    if any('image' in p for p in state['parts']):raise ValueError('the confidence layer reasoning confidence supports text only')

class ReasoningEngine(CudaEngine):
    def reasoning_identity(self):
        return {'base_runtime':self.identity,'reasoning_code_sha256':hashlib.sha256(Path(__file__).read_bytes()).hexdigest(),'adapter':'disabled','forced_suffix':FORCE}

    def estimate(self,state,block,n,stage='fast',max_new_tokens=512):
        if stage not in ('fast','reasoning'):raise ValueError('unknown stage')
        if any('image' in p for p in state['parts']):raise ValueError('the confidence layer text-only budget estimator')
        text=self._render(state['parts'],block)
        if stage=='reasoning':
            validate_request(state,n,max_new_tokens)
            marker='<|im_start|>assistant\n'
            if marker not in text:raise ValueError('unknown assistant template')
            text=text[:text.rfind(marker)]+marker+'<think>\n'
        ids=self.t.encode(text,add_special_tokens=False)
        prefix=state['prefix_ids'][0].tolist()
        if ids[:len(prefix)]!=prefix:raise ValueError('budget prefix mismatch')
        extra=len(self.t.encode(FORCE,add_special_tokens=False)) if stage=='reasoning' else 0
        return {'input_tokens':len(ids)-len(prefix)+extra,'generated_tokens':max_new_tokens if stage=='reasoning' else 0,'branches':1,'full_prompt_tokens':len(ids),'reused_prefix_tokens':len(prefix)}

    def _reason_inputs(self,state,block):
        text=self._render(state['parts'],block)
        marker='<|im_start|>assistant\n'
        if marker not in text:raise ValueError('unknown assistant template')
        text=text[:text.rfind(marker)]+marker+'<think>\n'
        ids,embeds,pos,_=self._encode(state['parts'],text,state['features'])
        if not self.torch.equal(ids[:,:state['prefix_tokens']],state['prefix_ids']):raise ValueError('reasoning prefix mismatch')
        return text,ids,embeds,pos

    def reason(self,state,block,n_letters,max_tokens=512,execution='cached',replay_tokens=None,cap_checkpoints=None):
        """Greedy base reasoning then bounded forced letter readout; replay is qualification only."""
        validate_request(state,n_letters,max_tokens)
        checkpoints=tuple(cap_checkpoints or ())
        if any(not isinstance(c,int) or isinstance(c,bool) or not 1<=c<=max_tokens for c in checkpoints):raise ValueError('invalid cap checkpoints')
        if checkpoints and (replay_tokens is not None or execution!='cached'):raise ValueError('checkpoints require cached generation')
        if self.placement!='question':raise ValueError('reasoning needs base document state')
        if execution not in ('cached','full'):raise ValueError('bad execution')
        torch=self.torch;p=state['prefix_tokens'];previous=self.ctx['start']
        from transformers import DynamicCache
        try:
            self.ctx['start']=None
            with torch.inference_mode():
                text,ids,embeds,pos=self._reason_inputs(state,block)
                cache=copy.deepcopy(state['cache']) if execution=='cached' else DynamicCache(config=self.lm.config)
                start=p if execution=='cached' else 0
                h=self.lm(inputs_embeds=embeds[:,start:],position_ids=pos[...,start:],past_key_values=cache,use_cache=True).last_hidden_state
                generated=[];snapshots={};stop='token_cap';lastpos=pos[...,-1:];close=self.t.encode('</think>',add_special_tokens=False)
                eos=self.t.eos_token_id
                for i in range(0 if execution=='full' and replay_tokens is not None else max_tokens):
                    token=int(replay_tokens[i]) if replay_tokens is not None and i<len(replay_tokens) else int(self.model.lm_head(h[:,-1]).argmax(-1))
                    if replay_tokens is not None and i>=len(replay_tokens):break
                    generated.append(token)
                    one=torch.tensor([[token]],device=ids.device);lastpos=lastpos+1
                    h=self.lm(input_ids=one,position_ids=lastpos,past_key_values=cache,use_cache=True).last_hidden_state
                    if len(generated) in checkpoints:
                        sf=self.t.encode(FORCE,add_special_tokens=False);sfids=torch.tensor([sf],device=ids.device)
                        sfpos=lastpos+torch.arange(1,len(sf)+1,device=ids.device).view(1,1,-1)
                        sh=self.lm(input_ids=sfids,position_ids=sfpos,past_key_values=copy.deepcopy(cache),use_cache=True).last_hidden_state
                        sl=self.model.lm_head(sh[:,-1]).float()[0]
                        li=[self.t.encode(c,add_special_tokens=False)[0] for c in 'ABCDEFGHIJ'[:n_letters]]
                        sc=sl[li];sp=sc.softmax(-1)
                        snapshots[str(len(generated))]={'letter_logits':sc.cpu().tolist(),'probabilities':sp.cpu().tolist(),'prediction':int(sp.argmax()),'mass':float(sl.softmax(-1)[li].sum()),'generation_tokens':len(generated),'finish_reason':'token_cap'}
                    if replay_tokens is None and (token==eos or generated[-len(close):]==close):stop='closed_reasoning' if token!=eos else 'eos';break
                if execution=='full' and replay_tokens is not None:
                    generated=list(replay_tokens[:max_tokens]);lastpos=pos[...,-1:]+len(generated)
                suffix=self.t.encode(FORCE,add_special_tokens=False);suffix_ids=torch.tensor([suffix],device=ids.device)
                suffix_pos=lastpos+torch.arange(1,len(suffix)+1,device=ids.device).view(1,1,-1)
                if execution=='full' and replay_tokens is not None:
                    extra=torch.tensor([generated+suffix],device=ids.device)
                    all_embeds=torch.cat((embeds,self.model.get_input_embeddings()(extra)),dim=1)
                    extra_pos=pos[...,-1:]+torch.arange(1,extra.shape[1]+1,device=ids.device).view(1,1,-1)
                    h=self.lm(inputs_embeds=all_embeds,position_ids=torch.cat((pos,extra_pos),dim=-1),use_cache=False).last_hidden_state
                else:
                    h=self.lm(input_ids=suffix_ids,position_ids=suffix_pos,past_key_values=cache,use_cache=True).last_hidden_state
                logits=self.model.lm_head(h[:,-1]).float()[0]
                letter_ids=[]
                for c in 'ABCDEFGHIJ'[:n_letters]:
                    token_ids=self.t.encode(c,add_special_tokens=False)
                    if len(token_ids)!=1:raise ValueError('non-single letter token')
                    letter_ids.append(token_ids[0])
                chosen=logits[letter_ids];probs=chosen.softmax(-1)
                return {'cap_checkpoints':snapshots,'letter_logits':chosen.cpu().tolist(),'probabilities':probs.cpu().tolist(),'mass':float(logits.softmax(-1)[letter_ids].sum()),'prediction':int(probs.argmax()),'generated_tokens':generated,'generation_tokens':len(generated),'text':self.t.decode(generated),'finish_reason':stop,'forced_answer':True,'execution':execution,'prompt_tokens':int(ids.shape[1]),'input_tokens':int(ids.shape[1])-p+len(suffix),'branch_tokens':int(ids.shape[1])-p+len(suffix),'reused_prefix_tokens':p if execution=='cached' else 0,'adapter_enabled':False,'identity':self.reasoning_identity()}
        finally:self.ctx['start']=previous

    def exit_features(self,state,block,n_letters,layers=(31,47,55)):
        """Full-depth feature capture only; explicitly does not execute early exit."""
        captured={};handles=[]
        def hook(layer):
            def capture(module,args,output):
                value=output[0] if isinstance(output,tuple) else output
                captured[layer]=value[:,-1].detach().clone()
            return capture
        try:
            for layer in layers:handles.append(self.lm.layers[layer].register_forward_hook(hook(layer)))
            full=self.ask(state,block,n_letters)
            result={}
            with self.torch.inference_mode():
                text=self._render(state['parts'],block);raw=self.t.encode(text,add_special_tokens=False)
                letter_ids=[self.t.encode(text+c,add_special_tokens=False)[-1] for c in 'ABCDEFGHIJ'[:n_letters]]
                for layer,value in captured.items():
                    logits=self.model.lm_head(self.lm.norm(value)).float()[0,letter_ids]
                    result[str(layer)]={'feature':value.float().cpu().tolist()[0],'letter_logits':logits.cpu().tolist(),'prediction':int(logits.argmax())}
            return {'full':{k:(v.tolist() if hasattr(v,'tolist') else v) for k,v in full.items()},'layers':result,'executed_truncated':False,'total_layers':len(self.lm.layers)}
        finally:
            for handle in handles:handle.remove()