"""Small, auditable byte-pair encoder trained only on supplied local text."""
import argparse, collections, json
from pathlib import Path
SPECIAL = ["<pad>", "<bos>", "<eos>", "<system>", "<user>", "<assistant>", "<tool>", "<plan>", "<memory>"]

class ByteBPETokenizer:
    def __init__(self, merges=None, special_tokens=SPECIAL):
        self.special_tokens = special_tokens
        self.special_to_id = {s:i for i,s in enumerate(special_tokens)}
        self.merges = [tuple(x) for x in (merges or [])] # child token IDs, in priority order
        self.base = len(special_tokens)
        self.pair_to_token = {p:self.base+256+i for i,p in enumerate(self.merges)}
        self.pair_rank = {p:i for i,p in enumerate(self.merges)}
        self.vocab_size = self.base + 256 + len(self.merges)
    def encode_bytes(self, raw):
        ids = [self.base+b for b in raw]
        # Standard BPE: always apply the currently lowest-ranked available pair.
        while len(ids) > 1:
            choices=[(i, self.pair_rank[(ids[i],ids[i+1])]) for i in range(len(ids)-1) if (ids[i],ids[i+1]) in self.pair_rank]
            if not choices: break
            i,rank=min(choices,key=lambda x:x[1]); ids[i:i+2]=[self.base+256+rank]
        return ids
    def encode(self, text, add_bos=False, add_eos=False):
        ids=[]; pos=0
        while pos < len(text):
            found=next(((s,i) for s,i in self.special_to_id.items() if text.startswith(s,pos)),None)
            if found: ids.append(found[1]);pos+=len(found[0]);continue
            ends=[text.find(s,pos) for s in self.special_to_id if text.find(s,pos)>=0]; end=min(ends) if ends else len(text)
            ids.extend(self.encode_bytes(text[pos:end].encode('utf8')));pos=end
        return ([self.special_to_id['<bos>']] if add_bos else [])+ids+([self.special_to_id['<eos>']] if add_eos else [])
    def decode(self, ids):
        def expand(x):
            if x < self.base: return b''
            if x < self.base+256: return bytes([x-self.base])
            a,b=self.merges[x-self.base-256]; return expand(a)+expand(b)
        out=[];data=bytearray()
        for x in ids:
            if x < self.base:
                if data: out.append(data.decode('utf8',errors='replace'));data.clear()
                out.append(self.special_tokens[x])
            else:data.extend(expand(x))
        if data:out.append(data.decode('utf8',errors='replace'))
        return ''.join(out)
    def save(self,path): Path(path).parent.mkdir(parents=True,exist_ok=True);Path(path).write_text(json.dumps({'merges':self.merges,'special_tokens':self.special_tokens}))
    @classmethod
    def load(cls,path):
        d=json.loads(Path(path).read_text());return cls(d['merges'],d['special_tokens'])

def train(files,vocab_size):
    symbols=[len(SPECIAL)+b for b in b''.join(Path(p).read_bytes() for p in files)];merges=[]
    while len(SPECIAL)+256+len(merges)<vocab_size:
        pairs=collections.Counter(zip(symbols,symbols[1:]));
        if not pairs:break
        pair,n=pairs.most_common(1)[0]
        if n<2:break
        token=len(SPECIAL)+256+len(merges);merges.append(pair);new=[];i=0
        while i<len(symbols):
            if i+1<len(symbols) and (symbols[i],symbols[i+1])==pair:new.append(token);i+=2
            else:new.append(symbols[i]);i+=1
        symbols=new
    return ByteBPETokenizer(merges)

def main():
 p=argparse.ArgumentParser(); sub=p.add_subparsers(dest='cmd',required=True); t=sub.add_parser('train'); t.add_argument('--input',required=True);t.add_argument('--out',required=True);t.add_argument('--vocab-size',type=int,default=16000);a=p.parse_args()
 files=list(Path(a.input).rglob('*.txt')); assert files, 'No .txt files found'; tok=train(files,a.vocab_size);tok.save(a.out);print(f'Saved {tok.vocab_size} tokens to {a.out}')
if __name__=='__main__': main()