"""Typed local inference adapter for the research decision checkpoints.

The response schema resembles TypeSafe's primitives. Confidence uses this
implementation's documented normalized maximum probability, not a claimed
reimplementation of TypeSafe's unpublished statistic. No text generation.
"""
from __future__ import annotations
import importlib.util
import math
from pathlib import Path


def prepare_runtime_profile(checkpoint, device='cuda:0'):
    # This profile is verified before any dependency import can choose kernels.
    import hashlib, json, sys
    root=Path(checkpoint);runtime=json.loads((root/'runtime.json').read_text())
    spec=runtime.get('normalization_profile')
    if spec is None:
        if '_decision_process_normalization_profile_v1' in sys.modules:
            raise RuntimeError('Use separate processes for profiled and unprofiled models')
        return None
    import torch
    target=torch.device(device)
    if target.type!='cuda' or not torch.cuda.is_available():
        raise RuntimeError('The bound profile requires a ROCm CUDA device')
    arch=getattr(torch.cuda.get_device_properties(target),'gcnArchName','').split(':')[0]
    if arch!=spec['validated_arch']:
        raise RuntimeError('Target GPU architecture does not match the bound profile: '+arch)
    relative=Path(spec['loader_file'])
    if relative.is_absolute() or '..' in relative.parts:raise ValueError('Unsafe profile loader path')
    path=root/relative
    if hashlib.sha256(path.read_bytes()).hexdigest()!=spec['loader_sha256']:
        raise ValueError('Bound runtime profile loader changed')
    definition=importlib.util.spec_from_file_location('decision_bundle_runtime_profile',path)
    module=importlib.util.module_from_spec(definition);definition.loader.exec_module(module)
    return module.ensure_profile(root)


def question_row(state, name, question):
    kind=question.get('type')
    if kind not in {'choice','noul','score'}:raise ValueError('Unknown question type')
    if 'instructions' not in question:raise ValueError('instructions is required')
    criteria=question.get('criteria')
    if kind=='noul':
        criteria={} if criteria is None else criteria
        if not isinstance(criteria,dict) or set(criteria)-{'true','false'}:
            raise ValueError('noul criteria may contain only true and false')
        options=[{'key':'false','description':criteria.get('false','The answer to the question is no.')},
                 {'key':'true','description':criteria.get('true','The answer to the question is yes.')}]
    elif kind=='score':
        if not isinstance(criteria,list) or not 2<=len(criteria)<=10:
            raise ValueError('score requires an ordered list of 2..10 criteria')
        options=[{'key':str(i),'description':value} for i,value in enumerate(criteria)]
    else:
        if not isinstance(criteria,dict) or not 2<=len(criteria)<=255:
            raise ValueError('choice requires a mapping of 2..255 criteria')
        if not all(isinstance(k,str) for k in criteria):raise ValueError('Choice keys must be strings')
        options=[{'key':key,'description':value} for key,value in criteria.items()]
    # The question name is used for bookkeeping only; encoders never render id.
    return {'id':name,'state':state,'instructions':question['instructions'],
            'options':options,'task_type':kind,'family':'inference'}


def typed_answer(row, probabilities):
    p=[float(v) for v in probabilities];k=len(row['options'])
    if len(p)!=k or any(not math.isfinite(v) or v<0 for v in p):
        raise ValueError('Invalid probability vector')
    total=sum(p)
    if total<=0 or abs(total-1)>1e-4:raise ValueError('Probabilities must sum to one')
    p=[v/total for v in p];selected=max(range(k),key=p.__getitem__)
    kind=row['task_type']
    if kind=='noul':
        keys=[o['key'] for o in row['options']]
        if set(keys)!={'false','true'}:raise ValueError('Native noul rows require false/true keys')
        return {'type':'noul','noul':p[keys.index('true')]}
    answer={'type':kind,'probabilities':{o['key']:v for o,v in zip(row['options'],p)},
            'confidence':max(0.,min(1.,(k*max(p)-1)/(k-1)))}
    if kind=='choice':answer['choice']=row['options'][selected]['key']
    else:
        if [o['key'] for o in row['options']] != [str(i) for i in range(k)]:
            raise ValueError('Native score rows require ordered numeric level keys')
        answer['score']=sum(i*v for i,v in enumerate(p))
        answer['legend']={str(i):o['description'] for i,o in enumerate(row['options'])}
    return answer


class DecisionEngine:
    def __init__(self, checkpoint, model_code, *, device='cuda:0', max_length=16384,
                 batch_size=8, temperatures=None, model_name='local-decision-research'):
        self.normalization_profile=prepare_runtime_profile(checkpoint, device=device)
        import torch
        path=Path(model_code)/'decision_model.py'
        spec=importlib.util.spec_from_file_location('research_decision_runtime',path)
        module=importlib.util.module_from_spec(spec);spec.loader.exec_module(module)
        model,tokenizer=module.DecisionModel.from_checkpoint(checkpoint,dtype=torch.bfloat16)
        self.model=model.to(device).eval();self.tokenizer=tokenizer;self.module=module
        self.device=device;self.max_length=max_length;self.batch_size=batch_size
        self.temperatures=temperatures or {};self.model_name=model_name
        if batch_size<1 or max_length<1:raise ValueError('Positive batch_size/max_length required')
        if any(not math.isfinite(v) or v<=0 for v in self.temperatures.values()):
            raise ValueError('Temperatures must be finite positive numbers')

    def predict_rows(self, rows):
        import torch
        encoded=[self.module.encode(row,self.tokenizer,self.max_length) for row in rows]
        pad=self.tokenizer.pad_token_id if self.tokenizer.pad_token_id is not None else self.tokenizer.eos_token_id
        records=[]
        with torch.inference_mode():
            for start in range(0,len(rows),self.batch_size):
                items=encoded[start:start+self.batch_size]
                batch={key:value.to(self.device) if torch.is_tensor(value) else value
                       for key,value in self.module.collate(items,pad).items()}
                with torch.autocast('cuda',dtype=torch.bfloat16):logits=self.model(**batch)
                for row,item,values in zip(rows[start:start+self.batch_size],items,logits):
                    k=len(row['options']);values=values[:k].float()
                    temperature=self.temperatures.get(row['task_type'],1.)
                    probabilities=(values/temperature).softmax(-1).tolist()
                    answer=typed_answer(row,probabilities)
                    prediction=max(range(k),key=probabilities.__getitem__)
                    if row['task_type']=='noul':
                        chosen='true' if answer['noul']>=.5 else 'false'
                        prediction=[o['key'] for o in row['options']].index(chosen)
                    rec={'id':row['id'],'status':'ok','prediction':prediction,
                         'probabilities':probabilities,'logits':values.tolist(),'temperature':temperature,
                         'native_contract':True,'truncated':False,'input_tokens':len(item['ids']),
                         'prompt_sha256':item['prompt_sha256'],'answer':answer}
                    if row['task_type']=='noul':rec['native_noul']=answer['noul']
                    if row['task_type']=='score':rec['native_score']=answer['score']
                    records.append(rec)
        return records

    def decide(self, state, questions):
        if not isinstance(questions,dict) or not questions:
            raise ValueError('questions must be a nonempty mapping')
        if not all(isinstance(name,str) for name in questions):raise ValueError('Question names must be strings')
        rows=[question_row(state,name,q) for name,q in questions.items()]
        result=self.predict_rows(rows)
        return {'model':self.model_name,'answers':{r['id']:r['answer'] for r in result},
                'usage':{'input_tokens':sum(r['input_tokens'] for r in result),'scored_questions':len(result)}}