"""
CyberSLM SFT — Prompt Formatter
================================
Converts normalised raw samples into tokenised ``(input_ids, labels)``
pairs ready for the data collator.

Tokenisation strategy (segment-based)
-------------------------------------
A sample is decomposed into an ordered list of ``(text, is_loss_target)``
segments (see ``ConversationTemplate`` for the canonical layout). Each
segment is encoded independently and the resulting token id lists are
concatenated to form ``input_ids``. ``labels`` mirror ``input_ids`` but with
non-target (prompt) positions set to ``IGNORE_INDEX`` (-100).

Why segment-based (and not char-offset) masking
------------------------------------------------
Locating the response boundary by re-encoding a prefix string and comparing
token counts is fragile: with ``add_dummy_prefix`` (and BPE merges across the
boundary) ``len(encode(prefix))`` need not equal the number of full-sequence
tokens covering that prefix, so the mask can drift by 1-2 tokens per turn.
Encoding each segment and concatenating gives an **exact** boundary because
``input_ids`` and ``labels`` are built from the very same token lists.

Special tokens (critical)
--------------------------
Special tokens are referenced by **id**, never as literal strings. The real
SentencePiece control ids are ``<bos>=2`` and ``<eos>=3``. Every assistant
response is terminated with ``eos_id`` as a genuine token and that ``eos_id``
is **left unmasked** so the model learns to stop. ``bos_id`` is prepended once
at the start of the sequence (masked). This replaces the previous, broken
approach of appending the literal characters ``"</s>"`` (which SentencePiece
encodes as ``<``,``/``,``s``,``>`` — never id 3), which prevented the model
from ever learning an end-of-sequence signal.
"""

from __future__ import annotations

import logging
from typing import List, Optional, Tuple

from configs.sft_config import SFTConfig, TemplateConfig

logger = logging.getLogger(__name__)

# Tokens with this label are excluded from the cross-entropy loss.
IGNORE_INDEX: int = -100

# One segment of the rendered prompt: (text, is_loss_target).
Segment = Tuple[str, bool]


# ---------------------------------------------------------------------------
# Segment builders (string level, EOS handled at token level downstream)
# ---------------------------------------------------------------------------

def build_alpaca_segments(
    instruction: str,
    output: str,
    input_text: str = "",
    tmpl: Optional[TemplateConfig] = None,
    include_response: bool = True,
) -> List[Segment]:
    """
    Build the ordered ``(text, is_loss_target)`` segments for an alpaca sample.

    The response text is a loss target; everything else (instruction, input,
    the response header) is masked. No EOS string is injected here — the
    end-of-sequence token is appended as a real token id by
    :func:`encode_segments`.
    """
    if tmpl is None:
        tmpl = TemplateConfig()

    if tmpl.strip_fields:
        instruction = instruction.strip()
        input_text = input_text.strip()
        output = output.strip()

    if input_text:
        prompt = (
            f"### Instruction:\n{instruction}\n\n"
            f"### Input:\n{input_text}\n\n"
            f"{tmpl.response_prefix}"
        )
    else:
        prompt = (
            f"### Instruction:\n{instruction}\n\n"
            f"{tmpl.response_prefix}"
        )

    segments: List[Segment] = [(prompt, False)]
    if include_response:
        segments.append((output, True))
    return segments


def build_conversation_segments(
    messages: List[dict],
    tmpl: Optional[TemplateConfig] = None,
    include_final_response: bool = True,
) -> List[Segment]:
    """
    Build ordered ``(text, is_loss_target)`` segments for a multi-turn
    conversation. Only assistant response bodies are loss targets.
    """
    if tmpl is None:
        tmpl = TemplateConfig()

    def _clean(s: str) -> str:
        return s.strip() if tmpl.strip_fields else s

    segments: List[Segment] = []

    # Optional leading system message.
    body = messages
    if messages and messages[0].get("role") == "system":
        sys_content = _clean(messages[0]["content"])
        if sys_content:
            segments.append((f"System: {sys_content}\n\n", False))
        body = messages[1:]

    i = 0
    n = len(body)
    while i < n:
        role = body[i]["role"]
        content = _clean(body[i]["content"])

        if role == "user":
            segments.append((f"### User:\n{content}\n\n", False))
            # Pair with the following assistant turn, if present.
            if i + 1 < n and body[i + 1]["role"] == "assistant":
                asst = _clean(body[i + 1]["content"])
                is_last = (i + 2 >= n)
                segments.append(("### Assistant:\n", False))
                if not (is_last and not include_final_response):
                    segments.append((asst, True))
                    segments.append(("\n\n", False))
                i += 2
            else:
                # Dangling user turn with no assistant reply -- the normal
                # generation case. The assistant header MUST still be emitted:
                # it is the cue the model was trained to continue from. Omitting
                # it (the previous behaviour) handed the model a prompt shaped
                # unlike anything in its training distribution.
                segments.append(("### Assistant:\n", False))
                i += 1

        elif role == "assistant":
            asst = content
            segments.append(("### Assistant:\n", False))
            segments.append((asst, True))
            segments.append(("\n\n", False))
            i += 1

        else:
            # Mid-conversation system message: treat as masked context.
            segments.append((f"### System:\n{content}\n\n", False))
            i += 1

    return segments


# ---------------------------------------------------------------------------
# Tokeniser wrapper
# ---------------------------------------------------------------------------

class Tokenizer:
    """
    Thin wrapper around a SentencePiece model that exposes only what the
    formatter and collator need.

    SentencePiece is not imported at module level so the formatter module can
    be imported in unit tests without requiring sentencepiece to be installed.
    """

    def __init__(self, model_path: str) -> None:
        try:
            import sentencepiece as spm  # type: ignore
        except ImportError as exc:
            raise ImportError(
                "sentencepiece is required for the tokenizer. "
                "Install with:  pip install sentencepiece"
            ) from exc

        self._sp = spm.SentencePieceProcessor()
        self._sp.Load(model_path)

        self.bos_id: int = self._sp.bos_id()
        self.eos_id: int = self._sp.eos_id()
        self.pad_id: int = self._sp.pad_id()
        self.vocab_size: int = self._sp.GetPieceSize()

        # Fail loudly if the tokenizer's control ids do not match the ids the
        # rest of the pipeline (and the base model config) assume. A silent
        # mismatch here corrupts training and breaks generation stopping.
        if self.eos_id < 0:
            raise ValueError(
                "Tokenizer has no EOS id (eos_id < 0). The SFT pipeline "
                "requires a real end-of-sequence token."
            )

    # ------------------------------------------------------------------
    def encode(
        self,
        text: str,
        add_bos: bool = False,
        add_eos: bool = False,
    ) -> List[int]:
        ids: List[int] = self._sp.Encode(text, out_type=int)
        if add_bos and self.bos_id >= 0:
            ids = [self.bos_id] + ids
        if add_eos and self.eos_id >= 0:
            ids = ids + [self.eos_id]
        return ids

    def decode(self, ids: List[int]) -> str:
        return self._sp.Decode(ids)

    def __len__(self) -> int:
        return self.vocab_size


# ---------------------------------------------------------------------------
# Segment → token id encoding
# ---------------------------------------------------------------------------

def encode_segments(
    segments: List[Segment],
    tokenizer: Tokenizer,
    max_seq_len: int,
    add_bos: bool = True,
) -> Tuple[List[int], List[int]]:
    """
    Encode ``(text, is_loss_target)`` segments into ``(input_ids, labels)``.

    * ``input_ids`` is the concatenation of each segment's token ids.
    * A real ``eos_id`` token is appended immediately after every loss-target
      (assistant) segment and is itself a loss target — the model must learn
      to emit it.
    * ``bos_id`` is prepended once at the start (masked) when ``add_bos``.
    * ``labels`` equal ``input_ids`` on target positions and ``IGNORE_INDEX``
      elsewhere.

    The sequence is truncated to ``max_seq_len`` tokens.
    """
    input_ids: List[int] = []
    labels: List[int] = []

    if add_bos and tokenizer.bos_id is not None and tokenizer.bos_id >= 0:
        input_ids.append(tokenizer.bos_id)
        labels.append(IGNORE_INDEX)

    for text, is_target in segments:
        ids = tokenizer.encode(text)
        input_ids.extend(ids)
        labels.extend(ids if is_target else [IGNORE_INDEX] * len(ids))
        if is_target:
            # Terminate the assistant turn with a learned EOS token.
            input_ids.append(tokenizer.eos_id)
            labels.append(tokenizer.eos_id)

    return input_ids[:max_seq_len], labels[:max_seq_len]


# ---------------------------------------------------------------------------
# Public API — PromptFormatter
# ---------------------------------------------------------------------------

class PromptFormatter:
    """
    Converts a normalised raw sample into a tokenised ``(input_ids, labels)``
    pair, applying loss-masking so only assistant responses are trained on.

    Parameters
    ----------
    cfg:
        Master ``SFTConfig`` (for template and data settings).
    tokenizer:
        Initialised ``Tokenizer`` instance.
    """

    def __init__(self, cfg: SFTConfig, tokenizer: Tokenizer) -> None:
        self.cfg = cfg
        self.tokenizer = tokenizer
        self.tmpl = cfg.template
        self.max_len = cfg.data.max_seq_len

    # ------------------------------------------------------------------
    def format(self, sample: dict) -> Optional[Tuple[List[int], List[int]]]:
        """
        Format a single normalised sample into ``(input_ids, labels)``.

        Returns ``None`` when the sample produces no trainable (non-masked)
        tokens — e.g. the response was truncated away.
        """
        if "messages" in sample:
            segments = build_conversation_segments(
                sample["messages"], tmpl=self.tmpl, include_final_response=True
            )
        else:
            segments = build_alpaca_segments(
                instruction=sample["instruction"],
                output=sample["output"],
                input_text=sample.get("input", ""),
                tmpl=self.tmpl,
                include_response=True,
            )

        # Nothing to train on if there are no target segments at all.
        if not any(is_target for _, is_target in segments):
            logger.debug("No assistant/response segments — skipping sample")
            return None

        input_ids, labels = encode_segments(
            segments, self.tokenizer, self.max_len, add_bos=True
        )

        if not input_ids or not any(l != IGNORE_INDEX for l in labels):
            logger.debug(
                "No response tokens remain after truncation at %d — skipping",
                self.max_len,
            )
            return None

        return input_ids, labels

    # ------------------------------------------------------------------
    def format_for_inference(self, sample: dict) -> List[int]:
        """
        Format a sample for *generation*: prompt only, no response body and no
        trailing EOS, with ``bos_id`` prepended.
        """
        if "messages" in sample:
            segments = build_conversation_segments(
                sample["messages"], tmpl=self.tmpl, include_final_response=False
            )
        else:
            segments = build_alpaca_segments(
                instruction=sample["instruction"],
                output=sample.get("output", ""),
                input_text=sample.get("input", ""),
                tmpl=self.tmpl,
                include_response=False,
            )
        input_ids, _ = encode_segments(
            segments, self.tokenizer, self.max_len, add_bos=True
        )
        return input_ids