"""Shared ChatML tokens, assistant supervision and context budgeting."""
import os
from functools import lru_cache
from pathlib import Path
from tokenizers import Tokenizer

ROOT = Path(__file__).resolve().parent.parent
SPECIALS = ('<|endoftext|>', '<|pad|>', '<|im_start|>', '<|im_end|>', '<|system|>', '<|user|>', '<|assistant|>', '<|reserved|>')
STOP_IDS = {0, 3}


@lru_cache(maxsize=1)
def tokenizer():
    return Tokenizer.from_file(str(ROOT / os.environ.get('VBQ_TOKENIZER', 'data/tokenizer/vbq-tokenizer.json')))


def render(messages, add_generation_prompt=False):
    ids, mask = [], []
    tok = tokenizer()
    for m in messages:
        role, content = m['role'], m['content']
        if role not in ('system', 'user', 'assistant') or not isinstance(content, str):
            raise ValueError('messages must contain system/user/assistant roles and string content')
        if any(s in content for s in SPECIALS):
            raise ValueError('message content cannot contain reserved tokens')
        header = tok.encode(f'<|im_start|>{role}\n', add_special_tokens=False).ids
        body = tok.encode(content, add_special_tokens=False).ids
        tail = tok.encode('\n', add_special_tokens=False).ids
        ids.extend(header + body + [3] + tail)
        mask.extend([0] * len(header) + [int(role == 'assistant')] * (len(body) + 1) + [0] * len(tail))
    if add_generation_prompt:
        header = tok.encode('<|im_start|>assistant\n', add_special_tokens=False).ids
        ids.extend(header); mask.extend([0] * len(header))
    return ids, mask


def fit_history(messages, budget):
    history = [dict(m) for m in messages]
    if not history or history[-1]['role'] != 'user':
        raise ValueError('chat must end with a user message')
    if any(m['role'] == 'system' for m in history[1:]):
        raise ValueError('system message must be first')
    start = int(history[0]['role'] == 'system')
    while True:
        ids, _ = render(history, add_generation_prompt=True)
        if len(ids) <= budget: return history, ids
        if len(history) <= start + 1:
            raise ValueError('system prompt and latest message exceed the context budget')
        end = start + 1
        while end < len(history) - 1 and history[end]['role'] != 'user': end += 1
        del history[start:end]


def stream_reply(model, messages, max_new=256, temperature=0.7, top_p=0.9, top_k=50, rep_penalty=1.1):
    import torch
    _, ids = fit_history(messages, model.c.max_seq_len - max_new)
    out, sent = [], ''
    for t in model.generate_cached(torch.tensor([ids], device=next(model.parameters()).device), max_new_tokens=max_new,
                                   temperature=temperature, top_p=top_p, top_k=top_k, repetition_penalty=rep_penalty, eos_id=STOP_IDS):
        if t in STOP_IDS: break
        out.append(t); text = tokenizer().decode(out)
        if text.endswith('�'): continue
        yield text[len(sent):]; sent = text
    text = tokenizer().decode(out)
    if len(text) > len(sent): yield text[len(sent):]