"""Reading Jais-ASR output: the leading `[...]` decision tag.

Every generation starts with a bracketed decision:

    "[EG] النص هنا"        ->  a transcript, in Egyptian Arabic
    "[msa] النص هنا"       ->  a transcript, in Modern Standard Arabic
    "[IFM-NoTranscript]"   ->  a deliberate abstention: no transcribable speech

The rule is general: the transcript never begins with a bracketed span. Anything the
model writes inside leading `[...]` spans is decision metadata, routed to `.tag` (and
to the API's `reasoning` field), never into `.text`:

  * a span whose body opens with 'IF' (any case) is the refusal branch: `.refused`
    is True and `.text` is empty -- nothing after a refusal is a transcript.
  * otherwise the first alphabetic component that case-folds onto a known tag sets
    `.language`; a span naming no supported language leaves `.language` as None.
  * no leading span -> `.tag` is None and `.text` is the whole output, untouched.

    from postprocess import split_text

    r = split_text(raw)
    r.text        # the transcript
    r.language    # canonical tag ("EG", "msa", ...) or None
    r.refused     # True when the model abstained
    r.tag         # the leading tag as emitted, e.g. "[EG]", or None

For well-formed output the exact token-id fast path (`split_token_ids`) gives the same
answer with no string work; the text rule is the authority otherwise.

This module is pure python: no torch, no transformers, no vllm at import time.
"""

from __future__ import annotations

import re
from dataclasses import dataclass

__all__ = [
    "LANGUAGE_TAGS",
    "REFUSAL_TAG",
    "Transcription",
    "split_text",
    "split_token_ids",
    "classify_tag_body",
    "stream_split",
    "verify_against_tokenizer",
]

# ---------------------------------------------------------------------------------
# The decision inventory. Token ids are those of this repo's tokenizer;
# `verify_against_tokenizer` below re-derives every one of them.
# ---------------------------------------------------------------------------------
BRACKET_OPEN_ID = 60  # '['
BRACKET_CLOSE_ID = 62  # ']'
THINK_OPEN_ID = 250029  # <ifm|think>
THINK_CLOSE_ID = 250030  # '</ifm|think>', the last token of every rendered prompt

REFUSAL_TAG = "[IFM-NoTranscript]"
REFUSAL_TOKEN_IDS = (
    60,
    5126,
    46,
    131126,
    195862,
    62,
)  # '[' 'IF' 'M' '-No' 'Transcript' ']'

#: decision-token id -> tag surface. 24 tags, all distinct.
LANGUAGE_TAG_IDS = {
    226606: "msa",  # Modern Standard Arabic
    272: "en",
    2075: "fr",
    273: "es",
    286: "it",
    7646: "EG",
    8048: "SA",
    6101: "MA",
    29354: "IQ",
    23129: "JO",
    20056: "SY",
    27526: "QA",
    64489: "DZ",
    23852: "LB",
    5594: "OM",
    28472: "YE",
    43852: "TN",
    17604: "MR",
    25893: "KW",
    9373: "SD",
    12933: "LY",
    5402: "AE",
    4295: "PS",
    41905: "BH",
}
LANGUAGE_TAGS = tuple(LANGUAGE_TAG_IDS.values())

#: longest well-formed prefix, in tokens -- the id fast path never looks further
MAX_TAG_TOKENS = len(REFUSAL_TOKEN_IDS)

#: a decision span is SHORT. 64 chars is far beyond any real tag; anything longer
#: is not tag-shaped and stays in the transcript untouched.
TAG_SPAN_MAX_CHARS = 64
#: total leading-metadata budget (a few stacked spans at most)
TAG_ZONE_MAX_CHARS = 160

# one bracketed span (no newline inside; the first ']' closes it); used with .match(pos)
_SPAN_RE = re.compile(r"[ \t]*\[([^\]\n]{0,%d})\]" % TAG_SPAN_MAX_CHARS)

_CANONICAL = {t.lower(): t for t in LANGUAGE_TAGS}


@dataclass(frozen=True)
class Transcription:
    """What the model decided, and what it transcribed."""

    text: str
    language: str | None = None
    refused: bool = False
    tag: str | None = None
    thinking: str = ""
    truncated: bool = False

    def __bool__(self) -> bool:  # `if result:` is False on an abstention
        return not self.refused and bool(self.text)


def has_prefilled_tag(prompt_ids):
    """Only a complete known [language] suffix counts as a prompt prefill."""
    return (
        len(prompt_ids) >= 3
        and int(prompt_ids[-3]) == BRACKET_OPEN_ID
        and int(prompt_ids[-2]) in LANGUAGE_TAG_IDS
        and int(prompt_ids[-1]) == BRACKET_CLOSE_ID
    )


def classify_tag_body(body: str):
    """(language | None, refused) for one span body. See the module contract."""
    b = body.strip()
    if b[:2].upper() == "IF":
        return None, True
    for comp in re.split(r"[^A-Za-z]+", b):
        if comp and comp.lower() in _CANONICAL:
            return _CANONICAL[comp.lower()], False
    return None, False


def _leading_spans(text: str):
    """Parse the leading metadata zone -> (spans, refused, language, end_pos)."""
    spans, refused, language, pos = [], False, None, 0
    while pos <= TAG_ZONE_MAX_CHARS:
        m = _SPAN_RE.match(text, pos)
        if not m:
            break
        body = m.group(1)
        spans.append(f"[{body}]")
        lang, ref = classify_tag_body(body)
        refused = refused or ref
        if language is None and lang is not None:
            language = lang
        pos = m.end()
    return spans, refused, language, pos


# ---------------------------------------------------------------------------------
# token-id path -- exact fast path for WELL-FORMED output
# ---------------------------------------------------------------------------------
def tag_span_len(token_ids) -> int:
    """Tokens spanned by a WELL-FORMED leading tag, else 0 (the text rule then decides)."""
    ids = list(token_ids[:MAX_TAG_TOKENS])
    if not ids or ids[0] != BRACKET_OPEN_ID:
        return 0
    if tuple(ids[: len(REFUSAL_TOKEN_IDS)]) == REFUSAL_TOKEN_IDS:
        return len(REFUSAL_TOKEN_IDS)
    if len(ids) >= 3 and ids[1] in LANGUAGE_TAG_IDS and ids[2] == BRACKET_CLOSE_ID:
        return 3
    return 0


def split_token_ids(token_ids, *, strip: bool = True):
    """-> (tag or None, refused, remaining_token_ids). Exact-form fast path only;
    anything else returns (None, False, ids) and callers fall back to `split_text`."""
    ids = list(token_ids)
    n = tag_span_len(ids)
    if n == 0:
        return None, False, ids
    if n == len(REFUSAL_TOKEN_IDS):
        return REFUSAL_TAG, True, ([] if strip else ids)
    return f"[{LANGUAGE_TAG_IDS[ids[1]]}]", False, (ids[n:] if strip else ids)


# ---------------------------------------------------------------------------------
# text path -- the authority
# ---------------------------------------------------------------------------------
def _split_think(text: str):
    """Generated-side think block (enable_thinking=True): the prompt opens
    '<ifm|think>', so the generation carries thought text up to the first
    '</ifm|think>'. Without that marker this is a no-op."""
    head = text.lstrip()
    if head.startswith("<ifm|think>"):
        text = head[len("<ifm|think>") :]
    idx = text.find("</ifm|think>")
    if idx < 0:
        return "", text
    return text[:idx].strip(), text[idx + len("</ifm|think>") :]


def split_text(text: str, *, strip: bool = True) -> Transcription:
    """Split a decoded generation per the general contract (module docstring)."""
    thinking, text = _split_think(text)
    spans, refused, language, pos = _leading_spans(text)
    if not spans:
        return Transcription(text=text, thinking=thinking)
    tag = " ".join(spans)
    if refused:
        return Transcription(
            text="" if strip else text, refused=True, tag=tag, thinking=thinking
        )
    rest = text[pos:].lstrip() if strip else text
    return Transcription(text=rest, language=language, tag=tag, thinking=thinking)


def stream_split(text: str):
    """Streaming view of the contract over the text-so-far.

    -> (reasoning_so_far, content_so_far, holding)

    * while the head could still be inside the metadata zone (only whitespace yet, or
      a short unclosed '[' span) the undecided tail is HELD (both strings frozen);
    * a refusal-branch output puts its closed leading span(s) in reasoning; nothing
      after a refusal is ever content;
    * otherwise reasoning is the closed leading span(s), content starts at the first
      non-metadata character. Both returned strings only ever grow (delta-safe).
    """
    spans, refused, _lang, pos = _leading_spans(text)
    if refused:
        return text[:pos], "", False
    rest = text[pos:]
    head = rest.lstrip(" \t")
    ws = len(rest) - len(head)
    if not spans:
        if head == "" and len(text) <= TAG_ZONE_MAX_CHARS:
            return "", "", True  # nothing decisive yet
        if (
            head.startswith("[")
            and "]" not in head
            and "\n" not in head
            and len(head) <= TAG_SPAN_MAX_CHARS + 2
        ):
            return "", "", True  # a span may be forming
        return "", text, False  # no metadata: all content
    if (
        head.startswith("[")
        and "]" not in head
        and "\n" not in head
        and len(head) <= TAG_SPAN_MAX_CHARS + 2
        and pos + ws <= TAG_ZONE_MAX_CHARS
    ):
        return text[:pos], "", True  # another span may be forming
    return text[:pos], text[pos + ws :], False


def verify_against_tokenizer(tokenizer) -> None:
    """Assert the id table matches this tokenizer. Cheap; call it in your tests."""
    enc = lambda s: tokenizer(s, add_special_tokens=False).input_ids  # noqa: E731
    assert (
        tuple(enc(REFUSAL_TAG)) == REFUSAL_TOKEN_IDS
    ), f"{REFUSAL_TAG} -> {enc(REFUSAL_TAG)}, expected {list(REFUSAL_TOKEN_IDS)}"
    assert tuple(enc("</ifm|think>")) == (
        THINK_CLOSE_ID,
    ), f"</ifm|think> -> {enc('</ifm|think>')}, expected [{THINK_CLOSE_ID}]"
    for tid, tag in LANGUAGE_TAG_IDS.items():
        ids = enc(f"[{tag}] ")
        assert (
            ids[0] == BRACKET_OPEN_ID and ids[1] == tid
        ), f"[{tag}] -> {ids[:3]}, expected decision id {tid} at position 1"
    assert len(set(LANGUAGE_TAG_IDS)) == 24


# =================================================================================
# vLLM integration -- metadata becomes `message.reasoning`, transcript is `content`
#
#   vllm serve <model> --trust-remote-code --reasoning-parser jais_asr_tag
#
# Registration happens on import inside a vLLM process; the vllm_plugin/ package in
# this repo does that automatically at server start.
# =================================================================================
try:  # pragma: no cover - only meaningful inside a vLLM process
    from vllm.reasoning.abs_reasoning_parsers import (
        ReasoningParser,
        ReasoningParserManager,
    )
    from vllm.entrypoints.openai.engine.protocol import DeltaMessage

    @ReasoningParserManager.register_module("jais_asr_tag")
    class JaisASRTagParser(ReasoningParser):
        """reasoning = the leading decision metadata; content = the transcript only.

        Same general contract as split_text: any leading bracketed span(s) go to
        `reasoning`; a refusal-branch output puts its leading tag in `reasoning`
        (the same value as `.tag`) with empty content; no leading span -> empty
        reasoning.
        """

        def __init__(self, tokenizer, *args, **kwargs):
            template = kwargs.get("chat_template_kwargs") or {}
            self.language_prefill = template.get("language_prefill")
            self._prefill_emitted = False
            super().__init__(tokenizer, *args, **kwargs)

        # ReasoningConfig auto-initialization tokenizes these at server start;
        # they are the prompt's think-block delimiters, already closed in every
        # rendered prompt, so the thinking-budget machinery stays inert.
        @property
        def reasoning_start_str(self) -> str:
            return "<ifm|think>"

        @property
        def reasoning_end_str(self) -> str:
            return "</ifm|think>"

        def is_reasoning_end(self, input_ids) -> bool:
            ids = list(input_ids)
            # vLLM also calls this with PROMPT ids (seeding per-request state) and
            # with prompt+output ids; every rendered prompt ends with '</ifm|think>'
            # (+ an optional prefill), so only what follows the LAST one is output.
            for i in range(len(ids) - 1, -1, -1):
                if ids[i] == THINK_CLOSE_ID:
                    ids = ids[i + 1 :]
                    break
            else:
                if ids and ids[0] == 0:  # BOS-led, no anchor: a bare prompt
                    return False
            text = self.model_tokenizer.decode(ids[:48])
            _r, content, holding = stream_split(text)
            return (not holding) and content != ""

        def extract_content_ids(self, input_ids):
            ids = list(input_ids)
            text = self.model_tokenizer.decode(ids)
            r = split_text(text)
            if r.refused:
                return []
            if r.tag is None:
                return ids
            # the space after ']' merges into the first content token, so keep the
            # boundary token: cut after the last prefix that fits INSIDE the tag zone
            cut = len(text) - len(r.text)
            keep = 0
            for k in range(1, len(ids) + 1):
                if len(self.model_tokenizer.decode(ids[:k])) > cut:
                    break
                keep = k
            return ids[keep:]

        def extract_reasoning(self, model_output, request):
            r = split_text(model_output)
            tag = r.tag
            if tag is None and not r.refused:
                kw = getattr(request, "chat_template_kwargs", None) or {}
                if kw.get("language_prefill"):
                    tag = f"[{kw['language_prefill']}]"  # tag was prefilled into the prompt
            reasoning = "\n".join(p for p in (r.thinking, tag or "") if p) or None
            if r.refused:
                return reasoning, ""
            if tag is None and not r.thinking:
                return None, model_output
            return reasoning, r.text

        def extract_reasoning_streaming(
            self,
            previous_text,
            current_text,
            delta_text,
            previous_token_ids,
            current_token_ids,
            delta_token_ids,
        ):
            pr, pc, _ = stream_split(previous_text)
            cr, cc, _holding = stream_split(current_text)
            dr = cr[len(pr) :] if cr.startswith(pr) else cr
            dc = cc[len(pc) :] if cc.startswith(pc) else cc
            if self.language_prefill and not self._prefill_emitted:
                dr = f"[{self.language_prefill}]" + ("\n" + dr if dr else "")
                self._prefill_emitted = True
            if not dr and not dc:
                return None
            return DeltaMessage(reasoning=dr or None, content=dc or None)

    __all__.append("JaisASRTagParser")
except (
    Exception
) as exc:  # vLLM absent, or an API shift -- the pure-python half still works
    import importlib.util

    if importlib.util.find_spec("vllm") is not None:
        import logging

        logging.getLogger(__name__).warning(
            "jais_asr_tag reasoning parser NOT registered (%s: %s); "
            "`--reasoning-parser jais_asr_tag` will fail with a KeyError",
            type(exc).__name__,
            exc,
        )