import torch
from .model import EXLLM,EXLLMConfig
from .tokenizer import HybridTokenizer

class UTF8State:
    __slots__=('need','lo','hi')
    def __init__(self): self.need=0; self.lo=0x80; self.hi=0xBF
    def accepts_byte(self,b):
        if self.need: return self.lo<=b<=self.hi
        return b in (9,10,13) or 0x20<=b<=0x7E or 0xC2<=b<=0xF4
    def push(self,b):
        if not self.accepts_byte(b): return False
        if self.need:
            self.need-=1; self.lo=0x80; self.hi=0xBF; return True
        if b<=0x7F:return True
        if 0xC2<=b<=0xDF:self.need=1
        elif b==0xE0:self.need=2;self.lo=0xA0
        elif 0xE1<=b<=0xEC or 0xEE<=b<=0xEF:self.need=2
        elif b==0xED:self.need=2;self.hi=0x9F
        elif b==0xF0:self.need=3;self.lo=0x90
        elif 0xF1<=b<=0xF3:self.need=3
        elif b==0xF4:self.need=3;self.hi=0x8F
        return True
    @property
    def complete(self):return self.need==0

def load_model(ckpt='weights/EXLLM-ONI5M-v1.0.pt'):
    tok=HybridTokenizer.load('tokenizer.json'); c=torch.load(ckpt,map_location='cpu',weights_only=False)
    m=EXLLM(EXLLMConfig(**c['config']));m.load_state_dict(c['model']);m.eval();return m,tok,c

@torch.inference_mode()
def generate(m,tok,prompt,max_new=56):
    ids=tok.encode_user(prompt); reserve=min(max_new,m.cfg.max_seq_len//2)
    if len(ids)>m.cfg.max_seq_len-reserve:
        keep=m.cfg.max_seq_len-reserve-3; body=ids[2:-1][-max(1,keep):];ids=[tok.BOS,tok.USER,*body,tok.ASSIST]
    out=[];state=UTF8State();lps=[]
    for _ in range(min(max_new,m.cfg.max_seq_len-len(ids))):
        x=torch.tensor([ids+out],dtype=torch.long);logits=m(x)[0,-1].clone()
        logits[tok.PAD]=logits[tok.BOS]=logits[tok.USER]=logits[tok.ASSIST]=-1e30
        valid=torch.zeros_like(logits,dtype=torch.bool)
        if state.complete:
            valid[tok.EOS]=True
            for t in range(256):
                if state.accepts_byte(t):valid[t]=True
            valid[256:256+len(tok.chars)]=True
        else:
            for t in range(256):
                if state.accepts_byte(t):valid[t]=True
        logits[~valid]=-1e30;lp=torch.log_softmax(logits,dim=-1);t=int(torch.argmax(logits));lps.append(float(lp[t]))
        if t==tok.EOS:break
        if t<256:state.push(t)
        elif not state.complete:continue
        out.append(t)
    while out:
        try:text=tok.decode(out);break
        except UnicodeDecodeError:out.pop()
    else:text=''
    return text.strip(),sum(lps)/max(1,len(lps))