"""Processor for Jais-ASR.

`processor(text=..., audio=..., sampling_rate=...)` takes one `<|ifm|audio|>`
placeholder per clip and expands it to `<audio>` + one placeholder per encoder frame
+ `</audio>`, so the model's `masked_scatter` lines up exactly with the frames the
audio tower emits.

Audio input is standardized automatically: pass `sampling_rate=` with whatever rate
your file has (a scalar, or one rate per clip) and it is resampled to the model's
16 kHz; multi-channel audio is downmixed to mono and integer PCM is converted to
float. When `sampling_rate` is omitted the audio is assumed to already be 16 kHz.
A bare 2-D numpy array is always a batch of mono clips (leading dim = batch); pass
a multi-channel clip inside a list to have its channels downmixed.

Audio length is unrestricted at the long end: no maximum duration, no truncation.
The mel is extracted per clip at full length and long audio is chunked inside the
encoder; each clip expands to its own frame count (~13/sec). Clips must be longer
than `n_fft // 2` samples (200 samples = 12.5 ms at 16 kHz), the STFT minimum.
"""

import re

import numpy as np

import transformers
from transformers.feature_extraction_utils import BatchFeature
from transformers.processing_utils import ProcessorMixin

if tuple(int(p) for p in transformers.__version__.split(".")[:2]) < (4, 57):
    raise ImportError(
        "MawJais-stt requires transformers >= 4.57 (earlier versions lack the "
        "WhisperFeatureExtractor attention-mask trim and corrupt the audio features); "
        f"found transformers=={transformers.__version__}. "
        "Upgrade with `pip install -U 'transformers>=4.57'`.")


def _audio_output_lengths(input_lengths):
    """Encoder output-frame count from the mel-frame count.

    Three stride-2 stages plus per-100-frame windowing: 100 mel frames/sec -> 13 output
    tokens/sec. Works elementwise on a python int or a torch / numpy integer tensor.
    """
    input_lengths_leave = input_lengths % 100
    feat_lengths = (input_lengths_leave - 1) // 2 + 1
    return ((feat_lengths - 1) // 2 + 1 - 1) // 2 + 1 + (input_lengths // 100) * 13


def _resample(wave, orig_sr, target_sr):
    """Resample one mono float32 clip with the best library installed."""
    try:
        import soxr
        return np.asarray(soxr.resample(wave, orig_sr, target_sr, quality="HQ"),
                          dtype=np.float32)
    except ImportError:
        pass
    try:
        import librosa
        return np.asarray(librosa.resample(wave, orig_sr=orig_sr, target_sr=target_sr),
                          dtype=np.float32)
    except ImportError:
        pass
    try:
        from math import gcd
        from scipy.signal import resample_poly
        g = gcd(int(orig_sr), int(target_sr))
        return np.asarray(resample_poly(wave, target_sr // g, orig_sr // g),
                          dtype=np.float32)
    except ImportError:
        pass
    try:
        import torch
        import torchaudio.functional as _taf
        return _taf.resample(torch.from_numpy(wave), orig_sr, target_sr).numpy()
    except ImportError:
        raise ImportError(
            f"this audio is {orig_sr} Hz and the model expects {target_sr} Hz; install a "
            "resampler (`pip install soxr`, or librosa / scipy / torchaudio) or resample "
            "the audio yourself.")


def _standardize_audio(wave, sr, target_sr):
    """One clip -> float32 mono at `target_sr`.

    Accepts any dtype (integer PCM is rescaled to [-1, 1]), mono or multi-channel
    (channels are averaged), any sample rate (`sr=None` means already `target_sr`).
    A clip that is already float mono at `target_sr` passes through untouched.
    """
    a = np.asarray(wave)
    # integer PCM -> [-1, 1] float FIRST: channel averaging below returns float and
    # would otherwise hide the integer dtype, leaving the waveform at PCM amplitude
    if np.issubdtype(a.dtype, np.unsignedinteger):      # e.g. uint8 PCM (offset binary)
        half = float(np.iinfo(a.dtype).max + 1) / 2.0
        a = (a.astype(np.float32) - half) / half
    elif np.issubdtype(a.dtype, np.integer):            # e.g. int16 / int32 PCM
        a = a.astype(np.float32) / float(np.iinfo(a.dtype).max + 1)
    if a.ndim == 2:
        ch_axis = 0 if a.shape[0] <= a.shape[1] else 1
        if a.shape[ch_axis] > 8:
            raise ValueError(
                f"cannot tell channels from samples in audio of shape {tuple(a.shape)}; "
                "pass each clip as 1-D mono or (channels, samples) / (samples, channels)")
        a = a.mean(axis=ch_axis)
    if a.ndim != 1:
        raise ValueError(
            f"each audio clip must be 1-D (mono) or 2-D (with a channel axis), "
            f"got shape {tuple(a.shape)}")
    a = a.astype(np.float32, copy=False)
    if sr is not None and int(sr) != int(target_sr):
        a = _resample(a, int(sr), int(target_sr))
    return a


class JaisASRProcessor(ProcessorMixin):
    attributes = ["feature_extractor", "tokenizer"]
    tokenizer_class = "AutoTokenizer"
    # transformers>=5 resolves the feature extractor from preprocessor_config.json
    # (`feature_extractor_type`) and deprecates the class attribute; <5 still needs it.
    if int(transformers.__version__.split(".")[0]) < 5:
        feature_extractor_class = "WhisperFeatureExtractor"
    valid_kwargs = ["chat_template", "audio_token", "boa_token", "eoa_token"]

    #: fields this processor can produce (read by serving backends)
    model_input_names = [
        "input_ids", "attention_mask", "input_features", "feature_attention_mask",
    ]

    def __init__(
        self,
        feature_extractor=None,
        tokenizer=None,
        chat_template=None,
        audio_token="<|ifm|audio|>",
        boa_token="<audio>",
        eoa_token="</audio>",
        **kwargs,
    ):
        self.audio_token = audio_token
        self.boa_token = boa_token       # span delimiters; never scattered onto
        self.eoa_token = eoa_token
        self.audio_token_id = (
            tokenizer.convert_tokens_to_ids(audio_token) if tokenizer is not None else None
        )
        super().__init__(feature_extractor, tokenizer, chat_template=chat_template)

    # ---- ids some serving backends look for by these exact names -----------------
    @property
    def audio_token_ids(self):
        return [self.audio_token_id]

    @property
    def audio_processor(self):
        return self.feature_extractor

    def __call__(
        self,
        text=None,
        audio=None,
        return_tensors=None,
        return_mm_token_type_ids: bool = False,
        **kwargs,
    ):
        if text is None:
            raise ValueError("`text` is required.")
        # rate of the INCOMING audio; consumed here (audio is standardized to the
        # feature extractor's rate below), never forwarded into the tokenizer
        sr_kw = kwargs.pop("sampling_rate", None)
        # ProcessorMixin.apply_chat_template(tokenize=True) always passes images= and
        # videos=; they must not fall through **kwargs into the tokenizer
        for key in ("images", "videos"):
            if kwargs.pop(key, None) is not None:
                raise ValueError(f"`{key}` is not supported: this is an audio-only model.")
        if isinstance(text, str):
            text = [text]

        audio_inputs = {}
        if audio is not None:
            sr_target = int(self.feature_extractor.sampling_rate)
            if isinstance(audio, np.ndarray):
                if audio.ndim == 1:
                    audio = [audio]      # one mono clip
                else:
                    audio = list(audio)  # leading dim is always the batch of clips
            elif not isinstance(audio, (list, tuple)):
                audio = [audio]
            if len(audio) == 0:
                raise ValueError("`audio` is empty; pass at least one clip or omit `audio=`.")
            if isinstance(sr_kw, (list, tuple)):
                if len(sr_kw) != len(audio):
                    raise ValueError(
                        f"sampling_rate has {len(sr_kw)} entries but {len(audio)} audio clips "
                        "were passed; give one rate per clip (or a single scalar).")
                audio = [_standardize_audio(a, s, sr_target) for a, s in zip(audio, sr_kw)]
            else:
                audio = [_standardize_audio(a, sr_kw, sr_target) for a in audio]
            min_samples = int(self.feature_extractor.n_fft) // 2 + 1
            for i, a in enumerate(audio):
                if len(a) < min_samples:
                    raise ValueError(
                        f"audio clip {i} has {len(a)} samples after standardization; the STFT "
                        f"needs at least {min_samples} samples "
                        f"({1000.0 * min_samples / sr_target:.1f} ms at {sr_target} Hz). Note a "
                        "bare 2-D array is a batch of mono clips; pass a multi-channel clip "
                        "inside a list.")
            n_placeholders = sum(t.count(self.audio_token) for t in text)
            if n_placeholders != len(audio):
                raise ValueError(
                    f"text contains {n_placeholders} '{self.audio_token}' placeholders but "
                    f"{len(audio)} audio clips were passed; they must match one-to-one (in order).")
            sr = self.feature_extractor.sampling_rate
            # each clip runs through the feature extractor ALONE: its mel values and
            # frame count never depend on which other clips share the batch
            feats = [
                np.asarray(
                    self.feature_extractor(
                        a,
                        sampling_rate=sr,
                        return_attention_mask=True,
                        padding="longest",
                        truncation=False,
                    )["input_features"]
                )[0]
                for a in audio
            ]
            mel_lens = [f.shape[-1] for f in feats]
            t_max = max(mel_lens)
            input_features = np.zeros((len(feats), feats[0].shape[0], t_max), dtype=np.float32)
            feature_attention_mask = np.zeros((len(feats), t_max), dtype=np.int32)
            for i, (f, n) in enumerate(zip(feats, mel_lens)):
                input_features[i, :, :n] = f
                feature_attention_mask[i, :n] = 1
            audio_inputs["input_features"] = input_features
            audio_inputs["feature_attention_mask"] = feature_attention_mask
            n_tokens = [int(_audio_output_lengths(int(n))) for n in mel_lens]
            repl = iter(
                f"{self.boa_token}{self.audio_token * n}{self.eoa_token}" for n in n_tokens
            )
            pattern = re.escape(self.audio_token)
            text = [re.sub(pattern, lambda _m: next(repl), t) for t in text]
        else:
            n_placeholders = sum(t.count(self.audio_token) for t in text)
            if n_placeholders:
                raise ValueError(
                    f"text contains {n_placeholders} '{self.audio_token}' placeholders but no "
                    "audio was passed; pass `audio=` with one clip per placeholder.")

        kwargs.setdefault("add_special_tokens", False)
        text_inputs = self.tokenizer(text=text, return_tensors=return_tensors, **kwargs)
        data = {**text_inputs, **audio_inputs}
        # PreTrainedTokenizerFast emits `token_type_ids` by default. This model has no
        # segment embeddings, and `generate()` rejects unknown kwargs outright, so
        # `model.generate(**processor(...))` would raise. Drop it here so the processor's
        # output is always directly usable.
        data.pop("token_type_ids", None)

        if return_mm_token_type_ids:
            # 1 wherever the audio embeddings will be scattered, 0 elsewhere. Mirrors
            # `input_ids` row for row (lists stay lists, so ragged batches work).
            import torch
            ids = data["input_ids"]
            if isinstance(ids, torch.Tensor):
                data["mm_token_type_ids"] = (ids == self.audio_token_id).long()
            elif isinstance(ids, np.ndarray):
                data["mm_token_type_ids"] = (ids == self.audio_token_id).astype(np.int64)
            else:
                data["mm_token_type_ids"] = [
                    [1 if t == self.audio_token_id else 0 for t in row] for row in ids
                ]

        return BatchFeature(data=data, tensor_type=return_tensors)

    def _get_num_multimodal_tokens(self, audio_lengths=None, **kwargs):
        """Placeholder count per clip, without running the feature extractor.

        `audio_lengths` are MEL-frame counts (100/sec). Used by serving backends to
        size the placeholder span before scheduling.
        """
        if audio_lengths is None or len(audio_lengths) == 0:
            return {"num_audio_tokens": []}
        return {"num_audio_tokens": [int(_audio_output_lengths(int(n))) for n in audio_lengths]}

    # ---- tokenizer passthroughs --------------------------------------------------
    def batch_decode(self, *args, **kwargs):
        return self.tokenizer.batch_decode(*args, **kwargs)

    def decode(self, *args, **kwargs):
        return self.tokenizer.decode(*args, **kwargs)

    # ---- reading the model's output ----------------------------------------------
    def post_process_transcription(self, generated, *, input_lengths=None, strip_tag=True,
                                   skip_special_tokens=True):
        """Decode generations and split off the leading `[...]` decision tag.

        `generated` is what `model.generate(...)` returned (ids, incl. the prompt — a
        batch or a single row) or already-decoded text. Pass `input_lengths` to trim
        the prompt first: either the padded prompt width as a single int, or the true
        prompt length per row (leading pad tokens are then skipped automatically, so
        per-row lengths stay correct under left padding).

        Returns a list of `Transcription(text, language, refused, tag, thinking,
        truncated)`; see postprocess.py. `.truncated` is True when a row hit the
        generation cap (no end-of-turn token in the generated ids -- raise
        `max_new_tokens` for long-form audio). With `strip_tag=False` the tag is
        left in `.text`.
        """
        try:
            from .postprocess import (split_text, BRACKET_OPEN_ID, BRACKET_CLOSE_ID,
                                      LANGUAGE_TAG_IDS, THINK_OPEN_ID, THINK_CLOSE_ID)
        except ImportError:
            if __package__:
                raise  # Preserve dependency errors raised inside sibling modules.
            from postprocess import (split_text, BRACKET_OPEN_ID, BRACKET_CLOSE_ID,
                                     LANGUAGE_TAG_IDS, THINK_OPEN_ID, THINK_CLOSE_ID)

        if isinstance(generated, str):
            return [split_text(generated, strip=strip_tag)]
        first = generated[0] if len(generated) else None
        if isinstance(first, str):
            return [split_text(t, strip=strip_tag) for t in generated]
        if getattr(generated, "ndim", None) == 1 or (
                first is not None and not hasattr(first, "__iter__")):
            generated = [generated]      # a single row of ids -> batch of one

        pad_id = self.tokenizer.pad_token_id
        bos_id = self.tokenizer.bos_token_id
        eos_ids = {self.tokenizer.eos_token_id,
                   self.tokenizer.convert_tokens_to_ids("<|ifm|endoftext|>")}
        eos_ids.discard(None)
        rows = []
        for i, seq in enumerate(generated):
            seq = [int(t) for t in seq]
            prefilled = None
            thinking_open = False
            if input_lengths is not None:
                if isinstance(input_lengths, int):
                    n = input_lengths          # padded prompt width (pads included)
                else:
                    n = int(input_lengths[i])  # true per-row prompt length: skip the
                    if pad_id is not None:     # leading pad run first (left padding)
                        lead = 0
                        while lead < len(seq) and seq[lead] == pad_id:
                            lead += 1
                        if lead and pad_id == bos_id:
                            lead -= 1          # that last one is the prompt's own BOS
                        n += lead
                if (n >= 3 and seq[n - 1] == BRACKET_CLOSE_ID
                        and seq[n - 3] == BRACKET_OPEN_ID
                        and seq[n - 2] in LANGUAGE_TAG_IDS):
                    prefilled = LANGUAGE_TAG_IDS[seq[n - 2]]  # tag prefilled into the prompt
                for token_id in reversed(seq[:n]):
                    if token_id == THINK_CLOSE_ID:
                        break
                    if token_id == THINK_OPEN_ID:
                        thinking_open = True
                        break
                seq = seq[n:]
            if seq and seq[0] == THINK_OPEN_ID:
                thinking_open = True
            pending_thinking = thinking_open and THINK_CLOSE_ID not in seq
            rows.append((self.tokenizer.decode(seq, skip_special_tokens=skip_special_tokens),
                         not any(t in eos_ids for t in seq), prefilled, pending_thinking))
        import dataclasses
        out = []
        for t, cut, pre, pending_thinking in rows:
            # A capped generation still inside <ifm|think> is reasoning, not a
            # transcript. The prompt establishes the open block when supplied.
            r = split_text(t + "</ifm|think>" if pending_thinking else t, strip=strip_tag)
            if r.tag is None and pre is not None:
                r = dataclasses.replace(r, tag=f"[{pre}]", language=pre)
            out.append(dataclasses.replace(r, truncated=cut))
        return out

    def language_pin(self, language, force_bracket=True):
        """A `LanguagePin` for this repo: constrain the language tag while keeping the
        model's own transcribe-or-abstain decision. `language` is a tag, an ISO code
        covering several tags (`"ar"`), or a list of tags.

            pin = processor.language_pin("ar")
            out = model.generate(**inputs,
                                 logits_processor=[pin.for_transformers(inputs["input_ids"].shape[1])])
        """
        try:
            from .language import LanguagePin
        except ImportError:
            if __package__:
                raise  # Preserve dependency errors raised inside sibling modules.
            from language import LanguagePin
        return LanguagePin(language, force_bracket=force_bracket)


__all__ = ["JaisASRProcessor"]