"""
chat_format.py -- single source of truth for the SFT / inference chat template.

After add_special_tokens.py grew the vocab, the model has ATOMIC tokens for
ChatML framing and reasoning/tool markers. This module builds training tensors
and inference prompts that use those exact tokens, so:

    * training (sft_*.py) and inference (infer_sft-chat.py) agree byte-for-byte, and
    * <|im_start|>, <|im_end|>, <think>, <begin_solution>, <tool_call>, ... are
      each a SINGLE token instead of being split into raw bytes.

Template (one turn):

    <|im_start|>{role}\n{content}<|im_end|>\n

Generation starts right after a trailing  "<|im_start|>assistant\n".
Loss is computed on assistant content + its closing <|im_end|> only; everything
else (system/user turns, headers) is masked with -100.
"""
from typing import List, Dict, Optional, Tuple

# Canonical atomic markers (must exist in tokenizer.json's special_tokens).
THINK_OPEN, THINK_CLOSE = "<think>", "</think>"
SOL_OPEN, SOL_CLOSE = "<begin_solution>", "<end_solution>"
IM_START, IM_END = "<|im_start|>", "<|im_end|>"

# Legacy / alternate marker spellings seen in older datasets -> canonical atomic
# tokens. Applied to every message's content so whatever convention the data
# uses collapses onto the tokens the model actually has.
MARKER_ALIASES = {
    "<|begin_of_thought|>": THINK_OPEN,
    "<|end_of_thought|>": THINK_CLOSE,
    "<|begin_of_solution|>": SOL_OPEN,
    "<|end_of_solution|>": SOL_CLOSE,
    # occasional variants
    "<thinking>": THINK_OPEN,
    "</thinking>": THINK_CLOSE,
}

_ROLE_MAP = {
    "human": "user", "user": "user", "prompter": "user",
    "gpt": "assistant", "assistant": "assistant", "bot": "assistant", "model": "assistant",
    "system": "system",
    "tool": "tool", "tool_response": "tool", "observation": "tool", "function": "tool",
}


def normalize_markers(text: str) -> str:
    if not text:
        return text
    for alias, canon in MARKER_ALIASES.items():
        if alias in text:
            text = text.replace(alias, canon)
    return text


def norm_role(msg: Dict) -> str:
    raw = msg.get("role") or msg.get("from") or msg.get("speaker") or "user"
    return _ROLE_MAP.get(str(raw).strip().lower(), "user")


def msg_content(msg: Dict) -> str:
    val = msg.get("content") if "content" in msg else msg.get("value")
    if val is None:
        val = msg.get("text", "")
    return val.strip() if isinstance(val, str) else str(val or "").strip()


def format_chat(history, system_prompt: Optional[str] = None,
                add_generation_prompt: bool = True) -> str:
    """Build an inference prompt string from (role, content) pairs (or dicts).
    Mirrors tokenize_chatml so inference matches training exactly."""
    s = ""
    if system_prompt:
        s += f"{IM_START}system\n{normalize_markers(system_prompt)}{IM_END}\n"
    for item in history:
        if isinstance(item, dict):
            role, content = norm_role(item), msg_content(item)
        else:
            role, content = item
            content = normalize_markers((content or "").strip())
        s += f"{IM_START}{role}\n{content}{IM_END}\n"
    if add_generation_prompt:
        s += f"{IM_START}assistant\n"
    return s


def stop_token_ids(tokenizer) -> List[int]:
    """Token ids that should halt generation: <|im_end|> (primary) + <eos>."""
    out = []
    try:
        imid = tokenizer.convert_tokens_to_ids(IM_END)
        if imid is not None and imid >= 0:
            out.append(int(imid))
    except Exception:
        pass
    eos = getattr(tokenizer, "eos_token_id", None)
    if eos is not None:
        out.append(int(eos))
    return sorted(set(out))


def tokenize_chatml(messages: List[Dict], tokenizer, max_length: int,
                    vocab_size: Optional[int] = None,
                    system: Optional[str] = None,
                    add_bos: bool = True) -> Optional[Tuple[List[int], List[int]]]:
    """
    Turn a message list into (input_ids, labels) with assistant-only loss.

    messages: list of dicts with role/from + content/value (any supported alias).
    system:   optional system prompt injected if the messages have none.
    Returns None if there is no trainable assistant content.
    """
    if not messages:
        return None
    enc = lambda t: tokenizer.encode(t, add_special_tokens=False)
    bos_id = getattr(tokenizer, "bos_token_id", None)

    msgs = list(messages)
    if system and not any(norm_role(m) == "system" for m in msgs if isinstance(m, dict)):
        msgs = [{"role": "system", "content": system}] + msgs

    input_ids, labels = [], []
    if add_bos and bos_id is not None:
        input_ids.append(bos_id)
        labels.append(-100)

    trained_any = False
    for m in msgs:
        if not isinstance(m, dict):
            continue
        role = norm_role(m)
        content = normalize_markers(msg_content(m))
        if not content:
            continue
        header = enc(f"{IM_START}{role}\n")
        body = enc(content)
        footer = enc(f"{IM_END}\n")
        input_ids += header
        labels += [-100] * len(header)
        if role == "assistant":
            input_ids += body + footer
            labels += body + footer          # train content + closing <|im_end|>
            trained_any = trained_any or bool(body)
        else:
            input_ids += body + footer
            labels += [-100] * (len(body) + len(footer))

    input_ids = input_ids[:max_length]
    labels = labels[:max_length]

    if vocab_size is not None:
        input_ids = [i if 0 <= i < vocab_size else 0 for i in input_ids]
        labels = [i if (i == -100 or 0 <= i < vocab_size) else -100 for i in labels]

    if not trained_any or not any(l != -100 for l in labels):
        return None
    return input_ids, labels