"""Reuse request-local gather indices for fixed-shape token-local operations."""
from collections import defaultdict
from .branch_attention import REQUEST
from contextvars import ContextVar
import types
import torch

DEPTH=ContextVar('spev_vectorized_token_ops_depth',default=0)


class PackingCache:
    def __init__(self):
        self.tree=None;self.device=None;self.plan=None

    def __call__(self,tree,device):
        if tree is self.tree and device==self.device:return self.plan
        buckets=defaultdict(list);start=tree.prefix_length
        for index,length in enumerate(tree.branch_lengths):
            buckets[((length+31)//32)*32].append((index,start,length));start+=length
        assert start==tree.length
        blocks=[];restore=[None]*len(tree.branch_lengths);offset=0
        for width,branches in sorted(buckets.items()):
            for begin in range(0,len(branches),8):
                block=branches[begin:begin+8];indices=[0]*(8*width);valid=[False]*(8*width)
                for slot,(index,start,length) in enumerate(block):
                    a=slot*width;indices[a:a+length]=range(start,start+length);valid[a:a+length]=[True]*length
                    restore[index]=torch.arange(length,device=device)+offset+a
                blocks.append((width,torch.tensor(indices,device=device,dtype=torch.long),torch.tensor(valid,device=device,dtype=torch.bool)))
                offset+=8*width
        order=torch.cat(restore) if restore else torch.empty(0,device=device,dtype=torch.long)
        self.tree=tree;self.device=device;self.plan=(blocks,order)
        return self.plan


def grouped_forward(original,tree_getter,plan_getter=None):
    cache=PackingCache() if plan_getter is None else plan_getter
    def forward(module,value,*args,**kwargs):
        if DEPTH.get():return original(value,*args,**kwargs)
        tree=tree_getter()
        if tree is None:raise ValueError('Grouped token operations require a tree layout')
        if value.ndim not in (2,3) or value.shape[-2]!=tree.length:raise ValueError('Grouped token operations require the packed token sequence')
        if value.ndim==3 and value.shape[0]!=1:raise ValueError('Grouped token operations support one packed request')
        x=value[0] if value.ndim==3 else value
        blocks,order=cache(tree,x.device);token=DEPTH.set(DEPTH.get()+1)
        try:
            prefix=original(x[:tree.prefix_length].contiguous(),*args,**kwargs);outputs=[]
            for width,indices,valid in blocks:
                packed=x.index_select(0,indices).masked_fill(~valid[:,None],0).reshape(8,width,x.shape[-1])
                output=original(packed,*args,**kwargs)
                assert isinstance(output,torch.Tensor) and output.shape[:2]==packed.shape[:2]
                outputs.append(output.reshape(8*width,output.shape[-1]))
            branch=(outputs[0] if len(outputs)==1 else torch.cat(outputs,dim=0)).index_select(0,order) if outputs else prefix[:0]
            output=torch.cat([prefix,branch],dim=0)
            return output[None] if value.ndim==3 else output
        finally:DEPTH.reset(token)
    return forward


def install(model,*,linears,moe):
    if moe:raise ValueError('MoE dispatch is not a token-local linear operation')
    originals=[]
    def cache(tree, device):
        request=REQUEST.get()
        attention_cache=request[1]
        if not hasattr(attention_cache, "packing"):
            attention_cache.packing=PackingCache()
        return attention_cache.packing(tree,device)
    getter=lambda:REQUEST.get()[0] if REQUEST.get() else None
    if linears:
        for layer in model.backbone.model.layers:
            for module in list(layer.modules()):
                is_lora=all(hasattr(module,name) for name in ['a','b','base','scale'])
                if isinstance(module,torch.nn.Linear) or is_lora:
                    original=module.forward;originals.append((module,original))
                    module.forward=types.MethodType(grouped_forward(original,getter,cache),module)
    def restore():
        for module,forward in reversed(originals):module.forward=forward
    return restore