"""Request-owned guarded hybrid cache for the fresh current-recipe graph."""
from contextlib import contextmanager
import torch
from transformers.cache_utils import Cache
from .runtime import execution_identity


def first_changed(a,b,dim=0):
    a=a.movedim(dim,0).contiguous();b=b.movedim(dim,0).contiguous()
    common=min(a.shape[0],b.shape[0])
    if not common:return 0
    av=a[:common].reshape(common,-1).view(torch.uint8)
    bv=b[:common].reshape(common,-1).view(torch.uint8)
    changed=(av!=bv).any(-1).nonzero().flatten()
    return int(changed[0]) if changed.numel() else common


def fingerprint(core):
    return tuple((n,p.data_ptr(),p._version,str(p.device),str(p.dtype),tuple(p.shape))
                 for n,p in list(core.named_parameters())+list(core.named_buffers()))


def execution_fingerprint(core):
    tensor=next(core.parameters(),None)
    if tensor is None:tensor=next(core.buffers(),None)
    return execution_identity(tensor.device if tensor is not None else 'cpu')


class CachedLinears:
    def __init__(self,base,cache):self.base=base;self.cache=cache
    def __getattr__(self,name):return getattr(self.base,name)
    def __call__(self,x,weight,bias=None,*,name,seed=0,bridge=False):
        rows=x.reshape(-1,x.shape[-1])
        key=(name,seed,bridge,weight.data_ptr(),weight._version,tuple(weight.shape),tuple(weight.stride()),
             x.dtype,str(x.device),None if bias is None else (bias.data_ptr(),bias._version))
        previous=self.cache.linears.get(key)
        start=0
        if previous is not None:
            old_input,old_output=previous
            changed=first_changed(rows,old_input)
            if changed==rows.shape[0]:
                value=old_output[:len(rows)]
                self.cache.stats['linear_hits']+=1
                self.cache.stats['reused_rows']+=len(rows)
                self.cache.linears[key]=(rows.detach().clone(),value.detach().clone())
                return value.reshape(*x.shape[:-1],weight.shape[0])
            start=(changed//128)*128
            self.cache.stats['reused_rows']+=start
            if changed<min(len(rows),len(old_input)):self.cache.stats['row_refreshes']+=1
        tail=self.base(rows[start:],weight,bias,name=name,seed=seed,bridge=bridge,activation_row_offset=start)
        value=tail if not start else torch.cat((previous[1][:start],tail),0)
        self.cache.stats['computed_rows']+=len(rows)-start
        self.cache.linears[key]=(rows.detach().clone(),value.detach().clone())
        return value.reshape(*x.shape[:-1],weight.shape[0])


class FidelHybridCache(Cache):
    """Projection histories and attention KV; Mamba scan replay is explicit."""
    def __init__(self,owner,core,prompt_length,context_limit=8192,mode='guarded'):
        super().__init__(layers=[])
        self.mode=mode;self.scans={}
        self.owner=owner;self.weight_fingerprint=fingerprint(core)
        self.execution_fingerprint=execution_fingerprint(core)
        self.prompt_length=prompt_length;self.context_limit=context_limit
        self.position=0;self.closed=False;self.token_ids=None
        self.linears={};self.attention={}
        self.stats={'linear_hits':0,'reused_rows':0,'computed_rows':0,'row_refreshes':0,
                    'attention_rows_reused':0,'attention_rows_computed':0,'decode_calls':0,
                    'scan_chunks_reused':0,'scan_chunks_computed':0,'scan_prefix_refreshes':0}
    def get_seq_length(self,layer_idx=0):return self.position
    def get_max_cache_shape(self,layer_idx=0):return self.context_limit
    def get_usable_length(self,new_seq_length,layer_idx=0):return self.position
    @property
    def is_compileable(self):return False
    def reorder_cache(self,beam_idx):
        if beam_idx.numel()!=1 or int(beam_idx[0])!=0:raise ValueError('Beam expansion is unsupported')
    def reset(self):
        self.linears.clear();self.attention.clear();self.scans.clear();self.token_ids=None
        self.position=0;self.closed=True
    def scan_for(self,index,live):
        def execute(q,k,v,**kwargs):
            from .mamba_scan import scan,ScanResult
            # The trapezoidal key scale at token t also reads dt/trap at t+1.
            # A state is committed only after this look-ahead token is present.
            boundary=((live-1)//16)*16
            prior=self.scans.get(index);start=0;state=None
            tensors={'q':(q,1),'k':(k,1),'v':(v,1),'z':(kwargs['z'],1),
                     'dt':(kwargs['dt'],2),'trap':(kwargs['trap'],2)}
            for name,value in kwargs['prepared'].items():
                if name=='angle_state':continue
                tensors['prepared/'+name]=(value,3 if name=='segsum' else (1 if name=='angles' else 2))
            if prior is not None and prior['boundary']:
                old_boundary=prior['boundary'];valid=old_boundary<=boundary
                for name,(value,axis) in tensors.items():
                    if name=='prepared/segsum':
                        valid=valid and first_changed(value,prior['inputs'][name],2)>=old_boundary//16
                    else:
                        needed=old_boundary+(1 if name in ('dt','trap') else 0)
                        valid=valid and first_changed(value,prior['inputs'][name],axis)>=needed
                    if not valid:break
                if valid:start=old_boundary;state=prior['state']
                else:self.stats['scan_prefix_refreshes']+=1
            result=scan(q,k,v,**kwargs,start=start,initial_state=state,checkpoint_at=boundary)
            output=result.output if not start else torch.cat((prior['output'][:,:start],result.output),1)
            self.stats['scan_chunks_reused']+=start//16
            self.stats['scan_chunks_computed']+=(q.shape[1]-start)//16
            self.scans[index]={'boundary':boundary,'state':result.checkpoint,'output':output.detach().clone(),
                               'inputs':{name:value.detach().clone() for name,(value,axis) in tensors.items()}}
            return ScanResult(output,result.state,result.last_key,result.trace,result.checkpoint,boundary)
        return execute

    def attend(self,q,k,v):
        prior=self.attention.get(7);start=0;outputs=[]
        if prior is not None:
            start=min(first_changed(q,prior['q'],2),first_changed(k,prior['k'],2),first_changed(v,prior['v'],2))
            start=min(start,q.shape[2])
            if start:outputs.append(prior['output'][:,:,:start])
        scale=q.shape[-1]**-.5
        for pos in range(start,q.shape[2]):
            from .operators import causal_attention_row
            outputs.append(causal_attention_row(q[:,:,pos:pos+1],k[:,:,:pos+1],v[:,:,:pos+1])[2])
        out=torch.cat(outputs,2)
        self.stats['attention_rows_reused']+=start
        self.stats['attention_rows_computed']+=q.shape[2]-start
        self.attention[7]={'q':q.detach().clone(),'k':k.detach().clone(),'v':v.detach().clone(),
                           'output':out.detach().clone()}
        return out


@contextmanager
def activated(core,cache):
    engine=core.engine;prior=getattr(core,'_active_cache',None)
    object.__setattr__(core,'engine',CachedLinears(engine,cache))
    object.__setattr__(core,'_active_cache',cache)
    try:yield
    finally:
        object.__setattr__(core,'engine',engine)
        object.__setattr__(core,'_active_cache',prior)


def cached_forward(core,new_ids,prompt_length,cache,owner,mode='guarded'):
    if cache is None:cache=FidelHybridCache(owner,core,prompt_length,mode=mode)
    if not isinstance(cache,FidelHybridCache) or cache.owner!=owner or cache.closed or cache.mode!=mode:
        raise ValueError('Use a live cache owned by this model')
    if cache.execution_fingerprint!=execution_fingerprint(core):
        cache.reset();raise ValueError('Execution backend changed; start a fresh request')
    if cache.prompt_length!=prompt_length or cache.weight_fingerprint!=fingerprint(core):
        cache.reset();raise ValueError('Prompt boundary or model tensors changed; start a fresh request')
    if new_ids.ndim!=2 or new_ids.shape[0]!=1 or new_ids.shape[1]<1:
        raise ValueError('Provide one or more new tokens for one request')
    full=new_ids if cache.token_ids is None else torch.cat((cache.token_ids,new_ids),1)
    if full.shape[1]>cache.context_limit:raise ValueError('Context limit exceeded')
    try:
        with activated(core,cache):logits=core.logits(full,prompt_length)[-1:].unsqueeze(0)
        cache.token_ids=full.detach().clone();cache.position=full.shape[1]
        cache.stats['decode_calls']+=1
        return logits,cache
    except BaseException:
        cache.reset();raise