"""Tokenizer surgeon (PRIVATE to Training_Module).

Safely registers <think>, </think> and <|im_end|> (+ <|im_start|>) as single tokens, makes sure the
embedding matrix has rows for them, and initialises new rows from the mean of their sub-token rows.
Also owns the shared chat format used by both training stages.
"""
from __future__ import annotations

import re
from typing import Dict, Iterable, List, Optional, Sequence

THINK_OPEN = "<think>"
THINK_CLOSE = "</think>"
IM_START = "<|im_start|>"
IM_END = "<|im_end|>"
ENDOFTEXT = "<|endoftext|>"
MODES = ("thinking", "non_thinking", "max_thinking")
DEFAULT_SYSTEM = "You are RSS AI, a helpful assistant created by RSS. You understand Hindi, Hinglish and English."


# Active vocabulary policy: ASCII (English/Hinglish/code/math/punctuation), Devanagari
# (including Vedic/extended marks), tokenizer control tokens, byte-fallback tokens, and
# every ID that actually occurs in the SFT examples. The original embedding table is
# deliberately NOT resized or rewritten: Stage-1 checkpoint keys/shapes stay unchanged.
_BYTE_FALLBACK = re.compile(r"^<0x[0-9A-Fa-f]{2}>$")
_SPECIAL_SPELLING = re.compile(r"^(?:<\|.*?\|>|<[^<>]+>)$")
_INDIC_RANGES = ((0x0900, 0x097F), (0x1CD0, 0x1CFF), (0xA8E0, 0xA8FF))
_INDIC_JOINERS = {0x200C, 0x200D}


def _is_active_token_spelling(token: str) -> bool:
    if not isinstance(token, str):
        return False
    if _BYTE_FALLBACK.match(token) or _SPECIAL_SPELLING.match(token):
        return True
    for char in token:
        code = ord(char)
        if code <= 0x7F or code in _INDIC_JOINERS:
            return True
        if any(lo <= code <= hi for lo, hi in _INDIC_RANGES):
            return True
    return False


def summarize_active_vocab(tokenizer, dataset_token_ids: Iterable[int], active_token_ids: Sequence[int],
                           vocab_rows: int) -> Dict[str, int]:
    """Count the retained token families and foreign-token pruning without guessing from ratios."""
    vocab = tokenizer.get_vocab()
    special = {int(x) for x in getattr(tokenizer, "all_special_ids", []) or [] if x is not None}
    for attr in ("bos_token_id", "eos_token_id", "pad_token_id", "unk_token_id", "sep_token_id",
                 "cls_token_id", "mask_token_id"):
        value = getattr(tokenizer, attr, None)
        if value is not None:
            special.add(int(value))
    categories = {"special_control": set(), "byte_fallback": set(),
                  "english_ascii_code_math": set(), "devanagari": set(), "foreign": set()}
    for spelling, raw_id in vocab.items():
        token_id = int(raw_id)
        if _BYTE_FALLBACK.match(spelling):
            categories["byte_fallback"].add(token_id)
        elif token_id in special or _SPECIAL_SPELLING.match(spelling):
            categories["special_control"].add(token_id)
        elif any(ord(ch) <= 0x7F for ch in spelling):
            categories["english_ascii_code_math"].add(token_id)
        elif any(ord(ch) in _INDIC_JOINERS or any(lo <= ord(ch) <= hi for lo, hi in _INDIC_RANGES)
                 for ch in spelling):
            categories["devanagari"].add(token_id)
        else:
            categories["foreign"].add(token_id)
    active = {int(x) for x in active_token_ids}
    data_only_foreign = categories["foreign"].intersection(active)
    return {"full_vocab_size": int(vocab_rows), "active_vocab_size": len(active),
            "inactive_vocab_size": int(vocab_rows) - len(active),
            "special_control_kept": len(categories["special_control"]),
            "byte_fallback_kept": len(categories["byte_fallback"]),
            "english_ascii_code_math_kept": len(categories["english_ascii_code_math"]),
            "devanagari_kept": len(categories["devanagari"]),
            "dataset_only_foreign_kept": len(data_only_foreign),
            "foreign_vocab_total": len(categories["foreign"]),
            "foreign_vocab_pruned": len(categories["foreign"] - data_only_foreign)}


def build_active_vocab_ids(tokenizer, dataset_token_ids: Iterable[int], vocab_rows: Optional[int] = None) -> List[int]:
    """Return sorted original token IDs permitted by the RSS Hindi/English/code policy.

    This is an output-vocabulary slice only. Input embeddings keep their original row IDs,
    preserving strict compatibility with an existing Stage-1 checkpoint. Dataset IDs are
    always unioned in, so no token present in an encoded SFT example can be pruned.
    """
    vocab = tokenizer.get_vocab()
    active = {int(token_id) for token_id in getattr(tokenizer, "all_special_ids", []) or []
              if token_id is not None and int(token_id) >= 0}
    for attr in ("bos_token_id", "eos_token_id", "pad_token_id", "unk_token_id",
                 "sep_token_id", "cls_token_id", "mask_token_id"):
        token_id = getattr(tokenizer, attr, None)
        if token_id is not None and int(token_id) >= 0:
            active.add(int(token_id))
    for spelling, token_id in vocab.items():
        if _is_active_token_spelling(spelling):
            active.add(int(token_id))
    active.update(int(token_id) for token_id in dataset_token_ids if int(token_id) >= 0)

    if vocab_rows is None:
        vocab_rows = max([int(v) for v in vocab.values()] + [-1]) + 1
    invalid = sorted(token_id for token_id in active if token_id >= int(vocab_rows))
    if invalid:
        raise ValueError(f"active tokenizer IDs exceed model vocabulary rows ({vocab_rows}): {invalid[:8]}")
    if not active:
        raise RuntimeError("active vocabulary is empty; refusing to construct a sliced loss")
    return sorted(active)


# ------------------------------------------------------------- chat format
def build_system(system: str, mode: str) -> str:
    return f"{(system or DEFAULT_SYSTEM).strip()}\n[MODE: {mode}]"


def render_prompt(system: str, user: str, mode: str) -> str:
    return (f"{IM_START}system\n{build_system(system, mode)}{IM_END}\n"
            f"{IM_START}user\n{user}{IM_END}\n"
            f"{IM_START}assistant\n")


def render_target(thinking: str, assistant: str, mode: str) -> str:
    if mode == "non_thinking":
        return f"{assistant}{IM_END}"
    return f"{THINK_OPEN}\n{thinking}\n{THINK_CLOSE}\n{assistant}{IM_END}"


# ----------------------------------------------------------------- surgery
def prepare_tokenizer(tokenizer, recipe_override: Optional[Dict] = None) -> Dict:
    """Register special tokens on the tokenizer. Returns the surgery info/recipe.

    recipe_override: {new_token_id: [sub-token ids]} saved by an earlier run. It is used when the tokens already
    exist in the tokenizer (so there is nothing to compute) but the weights still need the row initialisation."""
    vocab = tokenizer.get_vocab()
    wanted = [THINK_OPEN, THINK_CLOSE, IM_START, IM_END]
    missing = [t for t in wanted if t not in vocab]
    # sub-token recipe MUST be computed before the tokens exist as single tokens
    sources = {t: tokenizer.encode(t, add_special_tokens=False) for t in missing}
    if missing:
        existing: List[str] = list(getattr(tokenizer, "additional_special_tokens", []) or [])
        tokenizer.add_special_tokens({"additional_special_tokens": existing + missing})
    if tokenizer.pad_token is None:
        tokenizer.pad_token = ENDOFTEXT if ENDOFTEXT in tokenizer.get_vocab() else tokenizer.eos_token
    tokenizer.eos_token = IM_END
    tokenizer.model_max_length = 10 ** 9           # we chunk ourselves; silence length warnings
    tokenizer.padding_side = "right"

    recipe = {tokenizer.convert_tokens_to_ids(t): ids for t, ids in sources.items()}
    if not recipe and recipe_override:
        recipe = {int(k): [int(x) for x in v] for k, v in recipe_override.items()}
    eos_ids = [tokenizer.convert_tokens_to_ids(IM_END)]
    if ENDOFTEXT in tokenizer.get_vocab():
        eos_ids.append(tokenizer.convert_tokens_to_ids(ENDOFTEXT))
    return {
        "added": missing, "recipe": recipe, "resized": False,
        "ids": {"think_open": tokenizer.convert_tokens_to_ids(THINK_OPEN),
                "think_close": tokenizer.convert_tokens_to_ids(THINK_CLOSE),
                "im_end": tokenizer.convert_tokens_to_ids(IM_END),
                "im_start": tokenizer.convert_tokens_to_ids(IM_START),
                "pad": tokenizer.pad_token_id, "eos_ids": eos_ids},
    }


def finalize_model_embeddings(model, tokenizer, info: Dict) -> Dict:
    """Make sure the embedding matrix covers the tokenizer, init new rows, fix generation ids."""
    import torch

    needed = len(tokenizer)
    rows = int(getattr(model.config, "vocab_size", 0) or model.get_input_embeddings().weight.shape[0])
    ring = getattr(model, "_rss_ring", None)
    if needed > rows:
        if ring is not None:
            raise RuntimeError(f"tokenizer needs {needed} rows but the checkpoint only has {rows}; "
                               "resizing is not possible in tier-3 ring mode (use tier 2 once to resize).")
        model.resize_token_embeddings(needed, pad_to_multiple_of=8)
        info["resized"] = True
        print(f"   [tokenizer] embedding matrix resized {rows} -> {model.get_input_embeddings().weight.shape[0]}")
    else:
        print(f"   [tokenizer] {len(info['added'])} token(s) registered, embedding has {rows} rows -> no resize needed")

    if ring is None and info["recipe"]:
        with torch.no_grad():
            embed = model.get_input_embeddings().weight
            out = model.get_output_embeddings()
            out_w = out.weight if out is not None and out.weight.data_ptr() != embed.data_ptr() else None
            for new_id, src in info["recipe"].items():
                if not src:
                    continue
                idx = torch.tensor(src, dtype=torch.long, device=embed.device)
                embed[new_id] = embed[idx].float().mean(0).to(embed.dtype)
                if out_w is not None:
                    idx_o = torch.tensor(src, dtype=torch.long, device=out_w.device)
                    out_w[new_id] = out_w[idx_o].float().mean(0).to(out_w.dtype)
        print("   [tokenizer] new token rows initialised from the mean of their sub-token embeddings")

    gen = model.generation_config
    gen.eos_token_id = info["ids"]["eos_ids"]
    gen.pad_token_id = info["ids"]["pad"]
    return info