"""Choosing the transcription language in advance.

By default Jais-ASR decides for itself, at one generated position, between abstaining
(`[IFM-NoTranscript]`) and a language/dialect tag. This module lets a caller pin that
language WITHOUT taking away the model's ability to abstain.

Why that needs care
-------------------
The obvious implementation -- restrict the decision to {pinned_tag, IF} -- silently
changes the abstention rate: when the pinned tag is not the model's own first choice,
the refusal token can outscore it even on clear speech, and the model abstains where
it would have transcribed. So the two decisions are read separately, from the same
untouched logits:

    abstain  iff  logit(IF) > max over ALL tags      <- compare against all language tags
    else          argmax over the ALLOWED tags       <- the caller's pin

The abstain/transcribe decision is compared against the model's OWN best tag, never the
pinned one, so pinning does not change this comparison at a fixed decision prefix. Only which tag gets written changes.

Two ways to use it, with different guarantees:

  constrained  `LanguagePin` logits processor. Preserves the strict score comparison
               above (ties transcribe), not unconstrained sampling probabilities. Needs the
               processor to be passed to generate().
  prefill      write `[EG]` into the prompt (chat template `language_prefill=`). Works in any
               engine with no extra code, but the tag is then in the PROMPT, so
               `[IFM-NoTranscript]` can never be produced -- you always get text, even
               for silence. Use it when you want a transcript no matter what.
"""

from __future__ import annotations

import logging
import math

try:
    from .postprocess import (
        LANGUAGE_TAG_IDS,
        BRACKET_OPEN_ID,
        BRACKET_CLOSE_ID,
        has_prefilled_tag,
        THINK_OPEN_ID,
        THINK_CLOSE_ID,
    )
except ImportError:
    if __package__:
        raise
    from postprocess import (
        LANGUAGE_TAG_IDS,
        BRACKET_OPEN_ID,
        BRACKET_CLOSE_ID,
        has_prefilled_tag,
        THINK_OPEN_ID,
        THINK_CLOSE_ID,
    )

__all__ = [
    "SUPPORTED_LANGUAGES",
    "resolve_language",
    "LanguagePin",
    "prefill_for",
    "resolve_language_controls",
]

REFUSAL_DECISION_ID = 5126  # 'IF', the continuation of [IFM-NoTranscript]
TAG_TO_ID = {t: i for i, t in LANGUAGE_TAG_IDS.items()}

#: the 19 country codes plus MSA -- everything the model treats as Arabic
ARABIC_TAGS = (
    "msa",
    "EG",
    "SA",
    "MA",
    "IQ",
    "JO",
    "SY",
    "QA",
    "DZ",
    "LB",
    "OM",
    "YE",
    "TN",
    "MR",
    "KW",
    "SD",
    "LY",
    "AE",
    "PS",
    "BH",
)

#: ISO 639-1 -> the tags that count as that language. `ar` deliberately maps to ALL the
#: Arabic tags: pinning "Arabic" should let the model keep choosing the dialect.
ISO_TO_TAGS = {
    "ar": ARABIC_TAGS,
    "en": ("en",),
    "fr": ("fr",),
    "es": ("es",),
    "it": ("it",),
}

#: what vLLM's SupportsTranscription advertises on /v1/audio/transcriptions --
#: the ISO codes plus every concrete Jais-ASR tag, so dialect pinning is reachable
SUPPORTED_LANGUAGES = {
    "ar": "Arabic",
    "en": "English",
    "fr": "French",
    "es": "Spanish",
    "it": "Italian",
    "msa": "Arabic (Modern Standard)",
    **{t: f"Arabic ({t})" for t in ARABIC_TAGS if t != "msa"},
}


def resolve_language(language):
    """Normalise a caller's language spec into the set of allowed tags.

    Accepts an ISO 639-1 code ("ar"), a model tag ("EG", "msa"), a list of either, or
    None for "no restriction". Raises on anything unrecognised rather than silently
    falling back to auto -- a typo that quietly disables pinning is worse than an error.
    """
    if language is None:
        return None
    if not isinstance(language, (str, list, tuple)):
        raise ValueError(
            "language must be a tag, ISO code, non-empty list/tuple, or None"
        )
    specs = [language] if isinstance(language, str) else list(language)
    if not specs:
        raise ValueError("language list must not be empty; use None for no restriction")
    tags = []
    for s in specs:
        if not isinstance(s, str):
            raise ValueError("each language must be a string tag or ISO code")
        s = s.strip()
        if s in ISO_TO_TAGS:
            tags.extend(ISO_TO_TAGS[s])
        elif s in TAG_TO_ID:
            tags.append(s)
        elif s.lower() in ISO_TO_TAGS:
            tags.extend(ISO_TO_TAGS[s.lower()])
        else:
            raise ValueError(
                f"unknown language {s!r}. Use an ISO 639-1 code "
                f"({', '.join(sorted(ISO_TO_TAGS))}) or a model tag "
                f"({', '.join(sorted(TAG_TO_ID))})."
            )
    seen, out = set(), []
    for t in tags:
        if t not in seen:
            seen.add(t)
            out.append(t)
    return out


def prefill_for(language):
    """The literal prefix for the PREFILL path, e.g. 'EG' -> '[EG]'.

    Only a single concrete tag can be prefilled -- you cannot write a set into a prompt.
    Returns None when the spec resolves to more than one tag. The prefill ends at ']':
    the space before the first transcript word is the model's own to produce.
    """
    tags = resolve_language(language)
    if not tags or len(tags) != 1:
        return None
    return f"[{tags[0]}]"


class LanguagePin:
    """Pin the language while leaving the abstention decision untouched.

    Engine-agnostic: `apply(generated_ids, logits)` is the whole rule. Thin adapters for
    transformers (`__call__`) and for any engine that hands over (token_ids, logits).

        pin = LanguagePin("EG")                       # or "ar", or ["EG", "MA"]
        model.generate(**inputs, logits_processor=[pin.for_transformers(prompt_len)])
    """

    def __init__(self, language, force_bracket=True):
        self.tags = resolve_language(language)
        self.allowed_ids = [TAG_TO_ID[t] for t in self.tags] if self.tags else None
        self.all_tag_ids = list(LANGUAGE_TAG_IDS)
        self.force_bracket = force_bracket

    def apply(self, generated_ids, logits, *, thinking=False):
        """Constrain the decision zone in-place and return the score vector.

        Prompt tokens must be excluded. If the prompt opens a thinking block,
        pass thinking=True: scores are untouched until the generated closing
        marker, and the decision zone starts immediately after that marker.
        The adapters infer this flag from the prompt automatically.
        """
        if self.allowed_ids is None:
            return logits
        if thinking:
            for index, token_id in enumerate(generated_ids):
                if int(token_id) == THINK_CLOSE_ID:
                    generated_ids = generated_ids[index + 1 :]
                    break
            else:
                return logits
        n_generated = len(generated_ids)
        if n_generated > 2:
            return logits
        if n_generated == 0:
            return _keep_only(logits, BRACKET_OPEN_ID) if self.force_bracket else logits
        # With force_bracket=False, non-tag output must remain untouched.
        if int(generated_ids[0]) != BRACKET_OPEN_ID:
            return logits
        if n_generated == 2:
            if int(generated_ids[1]) in self.allowed_ids:
                return _keep_only(logits, BRACKET_CLOSE_ID)
            return logits  # Let [IFM-NoTranscript] continue naturally.

        # Read both decisions before masking. Refusal must strictly beat every tag;
        # ties transcribe. This is a deterministic score rule, not preservation of
        # the probability distribution of unconstrained sampling.
        best_tag = max(self.all_tag_ids, key=lambda i: float(logits[i]))
        abstains = float(logits[REFUSAL_DECISION_ID]) > float(logits[best_tag])
        keep_id = (
            REFUSAL_DECISION_ID
            if abstains
            else max(self.allowed_ids, key=lambda i: float(logits[i]))
        )
        return _keep_only(logits, keep_id)

    def for_transformers(self, prompt_len):
        """Bind to the padded prompt width; supports batch and expanded beam rows."""
        if not isinstance(prompt_len, int) or prompt_len <= 0:
            raise ValueError("prompt_len must be the positive padded prompt width")
        outer = self

        class _Proc:
            def __call__(self, input_ids, scores):
                if input_ids.shape[1] < prompt_len:
                    raise ValueError("input_ids is shorter than the bound prompt_len")
                if outer.allowed_ids is None:
                    return scores
                # Stateless normal-prompt fast path after the decision zone.
                if input_ids.shape[1] - prompt_len > 2:
                    if bool((input_ids[:, prompt_len - 1] == THINK_CLOSE_ID).all()):
                        return scores
                for b in range(scores.shape[0]):
                    prompt = input_ids[b, :prompt_len]
                    if has_prefilled_tag(prompt):
                        continue
                    outer.apply(
                        input_ids[b, prompt_len:],
                        scores[b],
                        thinking=has_open_thinking(prompt),
                    )
                return scores

        return _Proc()


def _keep_only(logits, token_id):
    keep = logits[token_id].clone() if hasattr(logits, "clone") else logits[token_id]
    if not math.isfinite(float(keep)):
        raise ValueError(
            "language pin selected a token without a finite score; check conflicting "
            "token constraints or invalid logits"
        )
    logits[:] = float("-inf")
    logits[token_id] = keep
    return logits


def has_open_thinking(prompt_ids):
    """Whether the final thinking delimiter in the prompt is an opening one."""
    # Tensor reversal copies the prompt and iteration unbinds every element.
    # Index backward so a final closing delimiter costs constant work.
    for index in range(len(prompt_ids) - 1, -1, -1):
        token_id = int(prompt_ids[index])
        if token_id == THINK_CLOSE_ID:
            return False
        if token_id == THINK_OPEN_ID:
            return True
    return False


#: The exact prompt contract, rendered here so engines that build their own prompt
#: (vLLM's transcription endpoint) do not have to run the jinja template. Kept in
#: exact agreement with chat_template.jinja: the system turn is rendered only when
#: there is system content.
GENERATION_PROMPT = (
    "<|ifm|begin_of_text|>"
    "{system_turn}"
    "<|ifm|im_start|>user\n{audio}<|ifm|im_end|>"
    "<|ifm|im_start|>assistant\n<ifm|think>\n</ifm|think>{prefill}"
)
SYSTEM_TURN = "<|ifm|im_start|>system\n{system}<|ifm|im_end|>"
AUDIO_TOKEN = "<|ifm|audio|>"


def resolve_language_controls(
    language=None, language_prefill=None, *, enable_thinking=False
):
    """Validate mutually exclusive pin/prefill options and canonicalize prefill."""
    if language is not None and language_prefill is not None:
        raise ValueError("Use either language (pinning) or language_prefill, not both")
    if language is not None:
        resolve_language(language)
    if language_prefill is not None:
        if not isinstance(language_prefill, str):
            raise ValueError("language_prefill must be one concrete language tag")
        tags = resolve_language(language_prefill)
        if len(tags) != 1:
            raise ValueError(
                "language_prefill must resolve to one tag; use language for a family"
            )
        if enable_thinking:
            raise ValueError(
                "language_prefill cannot be combined with enable_thinking=True"
            )
        language_prefill = tags[0]
    return language, language_prefill


def build_prompt(system="", language_prefill=None, audio_token=AUDIO_TOKEN):
    """Render the ASR prompt; only language_prefill inserts a decision tag."""
    _, tag = resolve_language_controls(language_prefill=language_prefill)
    return GENERATION_PROMPT.format(
        system_turn=SYSTEM_TURN.format(system=system) if system else "",
        audio=audio_token,
        prefill=f"[{tag}]" if tag else "",
    )