"""SentencePiece adapters for j72-tokenizer (reversible escape wrapper, escape-e000-v2)."""
from __future__ import annotations

class SentencePieceAdapter:
    kind="native_sentencepiece"
    def __init__(self, model):
        import sentencepiece as spm
        self.processor=spm.SentencePieceProcessor(model_file=str(model)); self.model=str(model)
    def encode(self,text): return self.processor.encode(text,out_type=int)
    def decode(self,ids): return self.processor.decode(ids)
    def pieces(self,ids): return [self.processor.id_to_piece(i) for i in ids]
    def is_byte(self,token):
        try:
            return bool(self.processor.is_byte(token))
        except (TypeError, ValueError):
            if isinstance(token, int):
                token = self.processor.id_to_piece(token)
            return bool(self.processor.is_byte(token))
    def is_unk(self,token): return token==self.processor.unk_id() or token=="<unk>"
    def unk_id(self): return self.processor.unk_id()
    def vocabulary(self): return [{"piece":self.processor.id_to_piece(i),"score":self.processor.get_score(i)} for i in range(self.processor.get_piece_size())]

class ReversibleSentencePieceAdapter(SentencePieceAdapter):
    """Escape wrapper kept separate from native SP results."""
    kind="reversible_sentencepiece_wrapper"
    ESC="\ue000"
    def _escape(self,text):
        return (text.replace(self.ESC,self.ESC+self.ESC).replace("\r\n",self.ESC+"R").replace("\r",self.ESC+"r").replace("\n",self.ESC+"n").replace("\u2028",self.ESC+"L").replace("\u2029",self.ESC+"P").replace("▁",self.ESC+"s").replace("\x00",self.ESC+"0"))
    def _unescape(self,text):
        out=[]; i=0
        while i<len(text):
            if text[i]!=self.ESC: out.append(text[i]); i+=1; continue
            if i+1>=len(text): raise ValueError("dangling sentinel")
            c=text[i+1]; out.append(self.ESC if c==self.ESC else {"R":"\r\n","r":"\r","n":"\n","L":"\u2028","P":"\u2029","s":"▁","0":"\x00"}.get(c,self.ESC+c)); i+=2
        return "".join(out)
    @classmethod
    def escape_text(cls, text):
        return cls._escape(object.__new__(cls), text)
    def encode(self,text): return self.processor.encode(self._escape(text),out_type=int)
    def decode(self,ids): return self._unescape(self.processor.decode(ids))