"""Preserve five-task trained-head confidence/routing with optional evidence."""
import hashlib,json
from pathlib import Path
import numpy as np
from solomon.engine import SolomonEngine,sha
from solomon.engine_reasoning import ReasoningEngine
from solomon.service_heads import HeadsService
from solomon.service_answers import Service as ConfidenceService
from solomon.service_evidence import EvidenceService,CONTRACT

class ServiceEngine(SolomonEngine,ReasoningEngine):
    def __init__(self,*args,correctness_report_sha256=None,**kwargs):
        super().__init__(*args,**kwargs)
        self.identity.update(trained_heads_report_sha256=correctness_report_sha256,serving_sha256=sha(__file__))
        self.identity.pop('fingerprint',None)
        self.identity['fingerprint']=hashlib.sha256(json.dumps(self.identity,sort_keys=True).encode()).hexdigest()

    def estimate(self,state,block,n,stage='fast',max_new_tokens=512):
        if not any('image' in p for p in state['parts']):
            return ReasoningEngine.estimate(self,state,block,n,stage=stage,max_new_tokens=max_new_tokens)
        if stage!='fast':raise ValueError('image reasoning is unqualified')
        with self._lock,self.torch.inference_mode():
            ids,_,_,_=self._encode(state['parts'],self._render(state['parts'],block),state['features'])
            p=state['prefix_tokens']
            if not self.torch.equal(ids[:,:p],state['prefix_ids']):raise ValueError('image estimate prefix mismatch')
            return {'input_tokens':int(ids.shape[1])-p,'generated_tokens':0,'branches':1,
                    'full_prompt_tokens':int(ids.shape[1]),'reused_prefix_tokens':p}

class ConfidenceHeadsService(HeadsService):
    def expected_identity(self,stage):
        return {**super().expected_identity(stage),'contract':CONTRACT,'serving_sha256':sha(__file__)}

    def _validate_artifact(self,stage,artifact):
        ConfidenceService._validate_artifact(self,stage,artifact)
        if artifact.get('schema')!='solomon-heads-envelope-v1':raise ValueError('the evidence layer confidence envelope required')
        selected=artifact.get('trained_head_identity',{});runtime=self.runtime_identity
        if (selected.get('weights_sha256')!=runtime.get('trained_heads_sha256') or
            selected.get('report_sha256')!=runtime.get('trained_heads_report_sha256') or
            selected.get('feature_identity',{}).get('runtime_fingerprint')!=runtime.get('base_fingerprint') or
            selected.get('feature_identity',{}).get('adapter_sha256')!=runtime.get('adapter_sha256')):
            raise ValueError('the evidence layer correctness/head/backbone identity mismatch')

class FreshDecider:
    def __init__(self,engine,design,store):self.engine=engine;self.design=design;self.store=store
    def __call__(self,parts,task,request):
        from solomon.service_states import Service as Readout
        engine=self.engine
        class Capture:
            def __init__(self):self.records=[]
            def ask(self,*args,**kwargs):
                result=engine.ask(*args,**kwargs);self.records.append(result);return result
        capture=Capture();state=engine.prefill(parts or [{'text':''}])
        class Fixed(Readout):
            def _warm(self,key):return state
        readout=Fixed(self.store,capture,self.design)
        with engine.task_context(task):answer=readout.ask('fresh-evidence',task,**request)
        return {'answer':answer,'cost':{'input_tokens':state['prefix_tokens']+sum(r['branch_tokens'] for r in capture.records),
                                       'branches':len(capture.records),'prefill_tokens':state['prefix_tokens']}}

def service(store,engine,design,*,artifact=None,views_artifact=None,views_tasks=(),selector=None,page_selector=None):
    answers=ConfidenceHeadsService(store,engine,design,artifact=artifact,views_artifact=views_artifact,views_tasks=views_tasks)
    return EvidenceService(answers,FreshDecider(engine,design,Path(store)/'.evidence-readout'),selector,page_selector)