import json, unicodedata
from pathlib import Path
BYTE_BASE=0; BYTE_COUNT=256

def normalize_text(s:str)->str:
    s=unicodedata.normalize('NFC',s)
    return ''.join((' ' if (ord(c)<32 and c not in '\n\t') else c) for c in s).strip()

class HybridTokenizer:
    def __init__(self, chars):
        self.chars=list(chars); self.char_to_id={c:256+i for i,c in enumerate(self.chars)}
        s=256+len(self.chars); self.PAD=s; self.BOS=s+1; self.USER=s+2; self.ASSIST=s+3; self.EOS=s+4; self.vocab_size=s+5
    @classmethod
    def load(cls,path): return cls(json.loads(Path(path).read_text(encoding='utf-8'))['chars'])
    def encode_text(self,s):
        out=[]
        for c in normalize_text(s):
            if c in self.char_to_id: out.append(self.char_to_id[c])
            else: out.extend(c.encode('utf-8','strict'))
        return out
    def encode_user(self,s): return [self.BOS,self.USER,*self.encode_text(s),self.ASSIST]
    def encode_example(self,prompt,answer,max_seq_len=128):
        p=self.encode_text(prompt); a=self.encode_text(answer)
        seq=[self.BOS,self.USER,*p,self.ASSIST,*a,self.EOS]
        if len(seq)>max_seq_len:
            excess=len(seq)-max_seq_len; p=p[min(excess,len(p)):]
            seq=[self.BOS,self.USER,*p,self.ASSIST,*a,self.EOS]
            if len(seq)>max_seq_len:
                a=a[:max(1,max_seq_len-(4+len(p)))]
                seq=[self.BOS,self.USER,*p,self.ASSIST,*a,self.EOS]
        return seq
    def token_bytes(self,tok):
        if 0<=tok<256: return bytes([tok])
        i=tok-256
        if 0<=i<len(self.chars): return self.chars[i].encode('utf-8')
        return b''
    def decode(self,toks):
        b=bytearray()
        for t in toks:
            if t in (self.PAD,self.BOS,self.USER,self.ASSIST,self.EOS): continue
            b.extend(self.token_bytes(t))
        return bytes(b).decode('utf-8','strict')