"""Native ByteLevel BPE tokenizer and explicit role-token chat formatting."""
from __future__ import annotations
from pathlib import Path

SPECIAL=['<|unk|>','<|pad|>','<|bos|>','<|eos|>','<|im_start|>','<|im_end|>','<|eot_id|>']
TOKENIZER_FILE='tokenizer.json'


def load_tokenizer():
    from tokenizers import Tokenizer
    from .settings import path
    p=path('tokenizer.json')
    if not p.exists():raise FileNotFoundError('Tokenizer missing: run python -m velqora tokenizer')
    return Tokenizer.from_file(str(p))


def train_tokenizer(corpus_file,vocab_size=32000):
    from tokenizers import Tokenizer,models,pre_tokenizers,decoders,trainers
    from .settings import path
    tokenizer=Tokenizer(models.BPE(unk_token='<|unk|>'))
    tokenizer.pre_tokenizer=pre_tokenizers.ByteLevel(add_prefix_space=False)
    tokenizer.decoder=decoders.ByteLevel()
    trainer=trainers.BpeTrainer(vocab_size=vocab_size,min_frequency=2,
                               initial_alphabet=pre_tokenizers.ByteLevel.alphabet(),
                               special_tokens=SPECIAL,show_progress=True)
    tokenizer.train([str(corpus_file)],trainer)
    tokenizer.save(str(path(TOKENIZER_FILE)))
    print(f'Saved tokenizer with {tokenizer.get_vocab_size()} tokens')
    return tokenizer


def require_ids(tokenizer):
    return {key: token_id(tokenizer,key) for key in SPECIAL}


def token_id(tokenizer,text):
    x=tokenizer.token_to_id(text)
    if x is None:raise ValueError(f'Missing special token: {text}')
    return x


def text_ids(tokenizer,text):
    return tokenizer.encode(text,add_special_tokens=False).ids


def chat_ids(tokenizer,messages,add_assistant_prefix=False,system=None):
    """Return (ids, loss_mask). Targets are assistant text/eot, not prompts or role markers."""
    spec=require_ids(tokenizer)
    start,end,eot=spec['<|im_start|>'],spec['<|im_end|>'],spec['<|eot_id|>']
    result=[];mask=[]
    if system is not None:
        messages=[{'role':'system','content':system}]+list(messages)
    for m in messages:
        role=m.get('role')
        content=m.get('content')
        if role not in ('user','assistant','system') or not isinstance(content,str):continue
        prefix=[start]+text_ids(tokenizer,role+'\n')
        content_tokens=text_ids(tokenizer,content)
        suffix=[eot if role=='assistant' else end]
        result.extend(prefix);mask.extend([0]*len(prefix))
        result.extend(content_tokens);mask.extend([1 if role=='assistant' else 0]*len(content_tokens))
        result.extend(suffix);mask.extend([1 if role=='assistant' else 0])
    if add_assistant_prefix:
        tail=[start]+text_ids(tokenizer,'assistant\n')
        result.extend(tail);mask.extend([0]*len(tail))
    return result,mask


def decode_reply(tokenizer,ids):
    spec=require_ids(tokenizer)
    clean=[]
    for i in ids:
        if i in (spec['<|eot_id|>'],spec['<|im_end|>'],spec['<|eos|>']):break
        clean.append(i)
    return tokenizer.decode(clean,skip_special_tokens=True)