"""Stable delimiter-based encoding used by the independent model."""
from __future__ import annotations

import re

QWEN_SPECIAL = ["<|fim_prefix|>", "<|fim_middle|>", "<|box_start|>", "<|box_end|>", "<|fim_suffix|>"]
# Reuse five reserved Gemma 4 control tokens instead of resizing its 262K
# vocabulary and training new embeddings. They are only structural markers in
# this prefill-only model; user text containing control-token syntax is escaped.
GEMMA4_SPECIAL = ["<|turn>", "<|channel>", "<|tool>", "<tool|>", "<turn|>"]
SPECIAL = QWEN_SPECIAL
SPECIAL_RE = re.compile(r"<\|([A-Za-z0-9_]+)\|>")


def user_tokens(tokenizer, text: str) -> list[int]:
    safe = SPECIAL_RE.sub(r"<¦\1¦>", str(text)).replace("<turn|>", "<turn¦>").replace("<tool|>", "<tool¦>")
    return tokenizer(safe, add_special_tokens=False).input_ids


def delimiter_tokens(tokenizer) -> list[str]:
    """Select five distinct, existing delimiter tokens for the backbone."""
    unknown = tokenizer.unk_token_id
    for profile in (QWEN_SPECIAL, GEMMA4_SPECIAL):
        ids = [tokenizer.convert_tokens_to_ids(token) for token in profile]
        if None not in ids and unknown not in ids and len(set(ids)) == len(ids):
            return profile
    raise ValueError("tokenizer has no supported five-token delimiter profile")


def _text_tokens(tokenizer, text: str) -> list[int]:
    return tokenizer(text, add_special_tokens=False).input_ids


def encode_question(tokenizer, record: dict, question: dict, max_state: int = 384, max_branch: int = 1024,
                    profile: str = "legacy") -> dict:
    if profile == "gemma_chat":
        state = user_tokens(tokenizer, record["state"])
        if len(state) + 1 > max_state:
            raise ValueError(f"state exceeds the configured {max_state}-token limit: {len(state) + 1}")
        bos = [tokenizer.bos_token_id] if tokenizer.bos_token_id is not None else []
        prefix = bos + _text_tokens(tokenizer, "<|turn>user\nContext:\n") + state
        instr = _text_tokens(tokenizer, "\nQuestion:\n") + user_tokens(tokenizer, question["instr"]) + _text_tokens(tokenizer, "\nOptions:\n")
        spans = [_text_tokens(tokenizer, "\nOption: ") + user_tokens(tokenizer, option) + _text_tokens(tokenizer, "\n")
                 for option in question["options"]]
        suffix = _text_tokens(tokenizer, "<turn|>\n<|turn>model\nDecision:")
        branch = instr + [token for span in spans for token in span] + suffix
        if len(branch) + len(prefix) > max_branch:
            raise ValueError("question branch exceeds the configured context")
        opt_ends, cursor = [], len(instr)
        for span in spans:
            cursor += len(span)
            opt_ends.append(cursor - 2 if len(span) > 1 else cursor - 1)
        return {"ids": prefix + branch, "state_len": len(prefix), "decide": len(prefix) + len(branch) - 1,
                "opts": [len(prefix) + index for index in opt_ends], "label": question.get("label", 0),
                "qtype": question.get("qtype", "choice"), "keys": question.get("keys", [])}
    if profile != "legacy":
        raise ValueError(f"unsupported encoding profile: {profile}")
    state = user_tokens(tokenizer, record["state"])
    if len(state) + 1 > max_state:
        raise ValueError(f"state exceeds the configured {max_state}-token limit: {len(state) + 1}")
    state_id, q_id, opt_id, end_id, decide_id = [tokenizer.convert_tokens_to_ids(x) for x in delimiter_tokens(tokenizer)]
    ids = [state_id] + state
    instr = [q_id] + user_tokens(tokenizer, question["instr"])
    spans = [[opt_id] + user_tokens(tokenizer, option) + [end_id] for option in question["options"]]
    branch = instr + [t for span in spans for t in span] + [decide_id]
    if len(branch) + len(ids) > max_branch:
        raise ValueError("question branch exceeds the configured context")
    opt_ends, cursor = [], len(instr)
    for span in spans:
        cursor += len(span)
        opt_ends.append(cursor - 1)
    return {"ids": ids + branch, "state_len": len(ids), "decide": len(ids) + len(branch) - 1,
            "opts": [len(ids) + i for i in opt_ends], "label": question.get("label", 0),
            "qtype": question.get("qtype", "choice"), "keys": question.get("keys", [])}