from pathlib import Path

import torch
import torchaudio
from torchaudio.compliance.kaldi import fbank
from transformers import ProcessorMixin
from transformers.feature_extraction_utils import BatchFeature


class SalmonnProcessor(ProcessorMixin):
    """Prepare text, audio, and contextual examples for SALMONN-2 inference."""

    attributes = ["tokenizer"]
    tokenizer_class = ("Qwen2Tokenizer", "Qwen2TokenizerFast")

    audio_placeholder = "<audio>"
    model_audio_placeholder = "<|vision_start|><|vision_end|>"

    def __init__(self, tokenizer, sample_rate=16000, num_mel_bins=128, chat_template=None):
        self.sample_rate = sample_rate
        self.num_mel_bins = num_mel_bins
        super().__init__(tokenizer, chat_template=chat_template)

    @property
    def model_input_names(self):
        return ["input_ids", "attention_mask", "audio_features", "audio_lengths", "audio_counts"]

    def prepare_model(self, model):
        """Register tokenizer-dependent timestamp tokens on a loaded SALMONN-2 model."""
        if getattr(model.config, "inject_temporal_embedding_nl", False):
            model.register_nl_timestamp_tokenizer(self.tokenizer)
        return model

    def build_prompt(self, instruction, audio_count=1, context=None):
        """Build the model prompt before chat templating.

        Context items may be strings for text-only contextual words or dictionaries
        containing ``text`` and, optionally, ``audio``. A context list must be
        consistently text-only or audio-text paired.
        """
        if not isinstance(instruction, str) or not instruction.strip():
            raise ValueError("instruction must be a non-empty string")
        if self.audio_placeholder in instruction:
            raise ValueError("instruction must not contain the reserved <audio> placeholder")
        if not isinstance(audio_count, int) or audio_count < 1:
            raise ValueError("audio_count must be a positive integer")

        prompt = self.audio_placeholder * audio_count + instruction.strip()
        contexts = self._normalize_context(context)
        if not contexts:
            return prompt

        has_audio = [item["audio"] is not None for item in contexts]
        if any(has_audio) and not all(has_audio):
            raise ValueError("context items must either all include audio or all be text-only")

        if all(has_audio):
            lines = [
                prompt,
                "Use the following contextual words and their pronunciations as references while transcribing the speech:",
                "<biasing_list>",
                *(f'{self.audio_placeholder}{item["text"]}' for item in contexts),
                "</biasing_list>.",
            ]
        else:
            words = ", ".join(item["text"] for item in contexts)
            lines = [
                prompt,
                "Pay extra attention to the following contextual words:",
                "<biasing_list>",
                f"[{words}]",
                "</biasing_list>.",
            ]
        return "\n".join(lines)

    def __call__(
        self,
        audios,
        instruction=None,
        context=None,
        formatted_prompt=None,
        sampling_rate=None,
        enable_thinking=False,
        return_tensors="pt",
        **tokenizer_kwargs,
    ):
        """Create a model-ready batch for one inference request.

        Args:
            audios: A path, waveform, or list of primary audio inputs.
            instruction: The user instruction associated with the primary audio. Use
                either ``instruction`` or ``formatted_prompt``, but not both.
            context: Optional contextual words. Each item is either a string or a
                ``{"text": ..., "audio": ...}`` dictionary.
            formatted_prompt: An advanced prompt containing one ``<audio>`` marker
                per input audio. This preserves custom audio placement and cannot be
                combined with ``instruction`` or ``context``.
            sampling_rate: Sampling rate for raw waveform inputs. Audio paths and
                ``{"array": ..., "sampling_rate": ...}`` inputs carry their own rate.
            enable_thinking: Whether to enable Qwen's thinking prompt.
            return_tensors: Currently only ``"pt"`` is supported.
        """
        if return_tensors != "pt":
            raise ValueError('SalmonnProcessor currently supports return_tensors="pt" only')

        primary_audios = self._as_audio_list(audios)
        if (instruction is None) == (formatted_prompt is None):
            raise ValueError("provide exactly one of instruction or formatted_prompt")

        if formatted_prompt is not None:
            if context is not None:
                raise ValueError("context cannot be combined with formatted_prompt")
            if not isinstance(formatted_prompt, str) or not formatted_prompt.strip():
                raise ValueError("formatted_prompt must be a non-empty string")
            if self.model_audio_placeholder in formatted_prompt:
                raise ValueError("formatted_prompt must use <audio>, not internal model audio markers")
            prompt = formatted_prompt.strip()
            ordered_audios = primary_audios
        else:
            contexts = self._normalize_context(context)
            context_audios = [item["audio"] for item in contexts if item["audio"] is not None]
            ordered_audios = primary_audios + context_audios
            prompt = self.build_prompt(instruction, len(primary_audios), contexts)

        expected_placeholders = len(ordered_audios)
        if prompt.count(self.audio_placeholder) != expected_placeholders:
            raise ValueError(
                f"Prompt contains {prompt.count(self.audio_placeholder)} audio placeholders "
                f"but received {expected_placeholders} audio inputs"
            )

        rendered = self.tokenizer.apply_chat_template(
            [{"role": "user", "content": prompt}],
            tokenize=False,
            add_generation_prompt=True,
            enable_thinking=enable_thinking,
        ).replace(self.audio_placeholder, self.model_audio_placeholder)

        tokenizer_kwargs.setdefault("add_special_tokens", False)
        encoded = self.tokenizer(rendered, return_tensors="pt", **tokenizer_kwargs)

        features = [self._extract_fbank(audio, sampling_rate) for audio in ordered_audios]
        audio_lengths = torch.tensor([item.size(0) for item in features], dtype=torch.long)
        audio_features = torch.nn.utils.rnn.pad_sequence(features, batch_first=True)
        audio_counts = torch.tensor([len(ordered_audios)], dtype=torch.long)

        return BatchFeature(
            data={
                **dict(encoded),
                "audio_features": audio_features,
                "audio_lengths": audio_lengths,
                "audio_counts": audio_counts,
            }
        )

    def decode(self, token_ids, clean_response=True, **kwargs):
        kwargs.setdefault("skip_special_tokens", True)
        text = self.tokenizer.decode(token_ids, **kwargs)
        return self._clean_response(text) if clean_response else text

    def batch_decode(self, sequences, clean_response=True, **kwargs):
        kwargs.setdefault("skip_special_tokens", True)
        texts = self.tokenizer.batch_decode(sequences, **kwargs)
        if clean_response:
            return [self._clean_response(text) for text in texts]
        return texts

    def _extract_fbank(self, audio, sampling_rate=None):
        waveform, source_rate = self._load_audio(audio, sampling_rate)
        if source_rate != self.sample_rate:
            waveform = torchaudio.functional.resample(waveform, source_rate, self.sample_rate)
        return fbank(
            waveform.unsqueeze(0),
            sample_frequency=self.sample_rate,
            num_mel_bins=self.num_mel_bins,
            low_freq=20.0,
            high_freq=-400.0,
            dither=0.0,
            snip_edges=False,
            energy_floor=1e-10,
        ).to(torch.float32)

    def _load_audio(self, audio, sampling_rate=None):
        if isinstance(audio, (str, Path)):
            waveform, source_rate = torchaudio.load(str(Path(audio).expanduser()))
        elif isinstance(audio, dict):
            if "array" not in audio or "sampling_rate" not in audio:
                raise ValueError("audio dictionaries require 'array' and 'sampling_rate' entries")
            waveform = torch.as_tensor(audio["array"])
            source_rate = audio["sampling_rate"]
        elif isinstance(audio, tuple) and len(audio) == 2:
            waveform = torch.as_tensor(audio[0])
            source_rate = audio[1]
        else:
            if sampling_rate is None:
                raise ValueError("sampling_rate is required for raw waveform inputs")
            waveform = torch.as_tensor(audio)
            source_rate = sampling_rate

        if waveform.ndim == 2:
            waveform = waveform.mean(dim=0)
        elif waveform.ndim != 1:
            raise ValueError("audio waveforms must have shape (samples,) or (channels, samples)")
        if waveform.numel() == 0:
            raise ValueError("audio waveforms must not be empty")
        if not isinstance(source_rate, int) or source_rate <= 0:
            raise ValueError("sampling_rate must be a positive integer")
        return waveform.to(torch.float32).cpu(), source_rate

    def _as_audio_list(self, audios):
        if isinstance(audios, list):
            if not audios:
                raise ValueError("audios must contain at least one audio input")
            return audios
        return [audios]

    def _normalize_context(self, context):
        if context is None:
            return []
        if not isinstance(context, (list, tuple)):
            raise TypeError("context must be a list of strings or dictionaries")

        normalized = []
        for item in context:
            if isinstance(item, str):
                text, audio = item, None
            elif isinstance(item, dict):
                text, audio = item.get("text"), item.get("audio")
            else:
                raise TypeError("each context item must be a string or dictionary")
            if not isinstance(text, str) or not text.strip():
                raise ValueError("each context item requires non-empty text")
            if self.audio_placeholder in text:
                raise ValueError("context text must not contain the reserved <audio> placeholder")
            normalized.append({"text": text.strip(), "audio": audio})
        return normalized

    @staticmethod
    def _clean_response(text):
        return text.replace("<think>", "").replace("</think>", "").strip()


__all__ = ["SalmonnProcessor"]