"""Chat template + loss masking for proto-buddy SFT.

Uses the role tokens already baked into the Tessera tokenizer:
  <|system|> <|user|> <|assistant|> <|end|>

Sequence layout (one training example):
  <|system|> {sys} <|end|> <|user|> {u} <|end|> <|assistant|> {a} <|end|>

Loss mask: ONLY the assistant content tokens + their closing <|end|> are training
targets. The model learns to PRODUCE replies (and to STOP at <|end|>), and never to
parrot the system/user turns. Everything else is context (label = ignore_index).
"""
from tokenizers import Tokenizer

DEFAULT_TOKENIZER = "tessera_tokenizer.json"
IGNORE = -100

# proto-buddy's turn is marked by a reserved slot designated as HIS self-token —
# never "<|assistant|>". The word "assistant" appears nowhere in his training.
BUDDY_TOK = "<|reserved_0|>"   # == "<|buddy|>" (cosmetic rename later if wanted)
# Injected-memory context lane: a foreign artifact (compact grounded notes), carried on its
# own reserved slot so the model learns to recognize "this is memory — use it carefully, may be
# stale/partial, is not the user's current words." It is CONTEXT, never a training target.
MEMORY_TOK = "<|reserved_1|>"  # == "<|buddy_memory|>"
ROLE_TOK = {"system": "<|system|>", "user": "<|user|>",
            "assistant": "<|assistant|>", "buddy": BUDDY_TOK, "memory": MEMORY_TOK}
TARGET_ROLES = ("assistant", "buddy")   # roles whose content is a training target
END = "<|end|>"


def load_tokenizer(path: str = DEFAULT_TOKENIZER) -> Tokenizer:
    return Tokenizer.from_file(path)


def _sid(tok: Tokenizer, s: str) -> int:
    i = tok.token_to_id(s)
    if i is None:
        raise KeyError(f"special token {s!r} not in tokenizer")
    return i


def build_example(tok: Tokenizer, messages: list[dict], max_len: int = 2048,
                  target_last_only: bool = False):
    """messages: [{role: system|user|buddy, content: str}, ...]
    Returns (ids, is_target) parallel lists. is_target[i] = a token the model should
    learn to emit. By default every buddy/assistant turn is a target; with
    target_last_only=True ONLY the final target-role turn is trained (the rest is
    context) — used for multi-turn corrections so we never reinforce earlier raw replies.
    """
    end = _sid(tok, END)
    ids: list[int] = []
    is_target: list[bool] = []
    last_tgt_idx = max((i for i, m in enumerate(messages) if m["role"] in TARGET_ROLES),
                       default=-1)

    def emit(token_ids, target):
        ids.extend(token_ids)
        is_target.extend([target] * len(token_ids))

    for i, m in enumerate(messages):
        role = m["role"]
        emit([_sid(tok, ROLE_TOK[role])], False)            # role cue: context, never a target
        content_ids = tok.encode(m["content"], add_special_tokens=False).ids
        train = role in TARGET_ROLES and (not target_last_only or i == last_tgt_idx)
        emit(content_ids, train)                             # reply -> TRAIN; else context
        emit([end], train)                                  #   + learn to STOP (only if trained)
    # LEFT-truncate (keep the TAIL): the trained turn — especially the only one under
    # target_last_only — sits at the END, so dropping the prefix preserves the target.
    # Right-truncation would silently yield an all-IGNORE example (no signal; NaN if a
    # whole batch hits it). Guard that the kept window still contains a target.
    if len(ids) > max_len:
        ids, is_target = ids[-max_len:], is_target[-max_len:]
        if not any(is_target):
            raise ValueError(f"example target turn exceeds max_len={max_len}; "
                             "raise max_len or shorten the final turn")
    return ids, is_target


def collate(batch, tok: Tokenizer, max_len: int = 2048, target_last_only: bool = False):
    """batch: list of messages-lists. Returns padded (input_ids, targets, attn-free).
    Shift convention matches ProtoGPT.forward (logits[t] predicts targets[t]):
      input_ids = ids[:-1],  targets[t] = ids[t+1] if assistant else IGNORE.
    """
    import torch
    pad = _sid(tok, "<|pad|>")
    rows = []
    for messages in batch:
        ids, tgt = build_example(tok, messages, max_len, target_last_only=target_last_only)
        inp = ids[:-1]
        lab = [ids[i + 1] if tgt[i + 1] else IGNORE for i in range(len(ids) - 1)]
        rows.append((inp, lab))
    L = max(len(r[0]) for r in rows)
    inp_b, lab_b = [], []
    for inp, lab in rows:
        padn = L - len(inp)
        inp_b.append(inp + [pad] * padn)
        lab_b.append(lab + [IGNORE] * padn)
    return torch.tensor(inp_b, dtype=torch.long), torch.tensor(lab_b, dtype=torch.long)


def render_prompt(tok: Tokenizer, messages: list[dict], gen_role: str = "buddy") -> list[int]:
    """Build a prompt ending at proto-buddy's turn token for generation (no reply yet)."""
    ids = []
    for m in messages:
        ids.append(_sid(tok, ROLE_TOK[m["role"]]))
        ids.extend(tok.encode(m["content"], add_special_tokens=False).ids)
        ids.append(_sid(tok, END))
    ids.append(_sid(tok, ROLE_TOK[gen_role]))
    return ids