"""AprielVoiceProcessor — Whisper audio, Gemma4 images, and text processing."""

import math
import re
from typing import Any, Dict, List, Optional, Tuple, Union

import numpy as np
import torch
import torch.nn.functional as F
import transformers

from .apriel_voice_config import AprielVoiceConfig

_SHORT_AUDIO_THRESHOLD_SECONDS = 2.0
_SHORT_AUDIO_SILENCE_SECONDS = 3.0


def _prepend_silence_to_short_audios(
    audios: List[np.ndarray], sampling_rate: int
) -> List[np.ndarray]:
    threshold_samples = int(_SHORT_AUDIO_THRESHOLD_SECONDS * sampling_rate)
    silence_samples = int(_SHORT_AUDIO_SILENCE_SECONDS * sampling_rate)

    return [
        np.pad(audio, (silence_samples, 0), mode="constant")
        if len(audio) < threshold_samples
        else audio
        for audio in audios
    ]


class AprielVoiceProcessor(transformers.ProcessorMixin):
    """Process audio and images for the AprielVoice multimodal model.

    Whisper features use Apriel's audio projector. Images use the wrapped
    Gemma4 vision tower, with one placeholder per actual image soft token.
    """

    attributes = ["audio_processor", "tokenizer", "image_processor"]
    audio_processor_class = ("WhisperFeatureExtractor",)
    tokenizer_class = ("PreTrainedTokenizer", "PreTrainedTokenizerFast")
    image_processor_class = ("Gemma4ImageProcessor", "Gemma4ImageProcessorPil")

    tokenizer: transformers.PreTrainedTokenizerBase
    audio_processor: transformers.FeatureExtractionMixin

    def __init__(
        self,
        audio_processor=None,
        tokenizer=None,
        audio_padding: str = "longest",
        encoder_ds_factor: int = 2,
        stack_factor: int = 8,
        audio_placeholder: str = "<|audio|>",
        audio_context_size: Optional[int] = 3000,
        image_processor=None,
    ):
        self.audio_padding = audio_padding
        self.encoder_ds_factor = encoder_ds_factor
        self.stack_factor = stack_factor
        self.audio_placeholder = audio_placeholder
        self.audio_context_size = audio_context_size
        assert tokenizer.eos_token is not None, "Tokenizer has no EOS token."
        self.vocab = tokenizer.get_vocab()
        self.audio_token_replacement = tokenizer.eos_token
        if tokenizer.pad_token_id is None:
            tokenizer.pad_token_id = tokenizer.eos_token_id
        if audio_processor is None:
            audio_processor = transformers.WhisperFeatureExtractor()
        if image_processor is None:
            image_processor = transformers.Gemma4ImageProcessor()
        self.image_token = getattr(tokenizer, "image_token", "<|image|>")
        self.image_token_id = self.vocab.get(self.image_token)
        self.boi_token = getattr(tokenizer, "boi_token", "<|image>")
        self.eoi_token = getattr(tokenizer, "eoi_token", "<image|>")
        super().__init__(
            audio_processor=audio_processor,
            tokenizer=tokenizer,
            image_processor=image_processor,
        )

    @classmethod
    def from_pretrained(cls, pretrained_model_name_or_path: str, **kwargs):
        image_kwargs = dict(kwargs.pop("images_kwargs", {}) or {})
        for key in transformers.Gemma4ImageProcessor.valid_kwargs.__annotations__:
            if key in kwargs:
                image_kwargs[key] = kwargs.pop(key)
        config: AprielVoiceConfig = transformers.AutoConfig.from_pretrained(
            pretrained_model_name_or_path, **kwargs
        )
        audio_processor = transformers.WhisperFeatureExtractor.from_pretrained(
            pretrained_model_name_or_path, **kwargs
        )
        tokenizer = transformers.AutoTokenizer.from_pretrained(
            pretrained_model_name_or_path, **kwargs
        )
        tokenizer.padding_side = "left"
        tokenizer.pad_token = tokenizer.eos_token
        vision_config = getattr(config.text_config, "vision_config", None)
        image_defaults = {}
        if vision_config is not None:
            image_defaults = {
                "patch_size": vision_config.patch_size,
                "pooling_kernel_size": vision_config.pooling_kernel_size,
                "max_soft_tokens": getattr(vision_config, "default_output_length", 280),
            }
        image_processor = transformers.Gemma4ImageProcessor(
            **{**image_defaults, **image_kwargs}
        )
        return cls(
            audio_processor=audio_processor,
            tokenizer=tokenizer,
            stack_factor=config.stack_factor,
            image_processor=image_processor,
        )

    def _auto_chunk_audios(
        self,
        text: str,
        audios: List[np.ndarray],
        sampling_rate: int,
        max_chunk_duration: float = 30.0,
    ) -> Tuple[str, List[np.ndarray]]:
        max_samples = int(max_chunk_duration * sampling_rate)
        new_audios: List[np.ndarray] = []
        chunk_counts: List[int] = []

        for audio in audios:
            if len(audio) > max_samples:
                n_chunks = math.ceil(len(audio) / max_samples)
                for i in range(n_chunks):
                    new_audios.append(audio[i * max_samples : (i + 1) * max_samples])
                chunk_counts.append(n_chunks)
            else:
                new_audios.append(audio)
                chunk_counts.append(1)

        if all(c == 1 for c in chunk_counts):
            return text, audios

        audio_block_re = re.compile(
            r"\[AUDIO FILE \d+ START\]\s*<\|audio>\s*<\|audio\|>\s*<audio\|>\s*\[AUDIO FILE \d+ END\]"
        )
        blocks = list(audio_block_re.finditer(text))

        if len(blocks) == len(chunk_counts):
            new_text_parts: List[str] = []
            prev_end = 0
            file_idx = 1
            for block, count in zip(blocks, chunk_counts):
                new_text_parts.append(text[prev_end:block.start()])
                for j in range(count):
                    if j > 0:
                        new_text_parts.append(" ")
                    new_text_parts.append(
                        f"[AUDIO FILE {file_idx} START] <|audio> <|audio|> <audio|> [AUDIO FILE {file_idx} END]"
                    )
                    file_idx += 1
                prev_end = block.end()
            new_text_parts.append(text[prev_end:])
            text = "".join(new_text_parts)
        else:
            parts = text.split("<|audio|>")
            new_parts: List[str] = []
            audio_idx = 0
            for i, part in enumerate(parts):
                new_parts.append(part)
                if i < len(parts) - 1:
                    count = chunk_counts[audio_idx] if audio_idx < len(chunk_counts) else 1
                    new_parts.append("<|audio|>" * count)
                    audio_idx += 1
            text = "".join(new_parts)

        return text, new_audios

    def _chunk_and_pad_audio(
        self,
        audio_values: torch.Tensor,
        audio_lens: torch.Tensor,
        include_audio_num_chunks: bool = False,
    ) -> Dict[str, Any]:
        chunked_audio_values: List[torch.Tensor] = []
        chunked_audio_lens: List[int] = []
        is_continuation_list: List[bool] = []
        num_chunks: List[int] = []
        context_size = self.audio_context_size or audio_values.shape[-1]

        for i in range(audio_values.shape[0]):
            num_chunks.append((int(audio_lens[i]) + context_size - 1) // context_size)
            for offset in range(0, audio_lens[i], context_size):
                is_continuation = offset > 0
                chunk = audio_values[i, :, offset : offset + context_size]
                if is_continuation and chunk.shape[-1] < context_size:
                    chunk = F.pad(chunk, (0, context_size - chunk.shape[-1]))
                chunked_audio_values.append(chunk)
                chunked_audio_lens.append(min(int(audio_lens[i].item()) - offset, context_size))
                is_continuation_list.append(is_continuation)

        data = {
            "audio_values": torch.stack(chunked_audio_values, dim=0),
            "audio_lens": torch.tensor(
                chunked_audio_lens, dtype=torch.int64, device=audio_values.device
            ),
            "audio_is_continuation": torch.tensor(
                is_continuation_list, dtype=torch.bool, device=audio_values.device
            ),
            "audio_batch_size": torch.tensor(
                [len(chunked_audio_values)], device=audio_values.device
            ),
        }
        if include_audio_num_chunks:
            data["audio_num_chunks"] = torch.tensor(
                num_chunks, dtype=torch.int64, device=audio_values.device
            )
        return data

    def __call__(
        self,
        text: Optional[str] = None,
        audio: Optional[Union[np.ndarray, torch.Tensor]] = None,
        audios: Optional[
            Union[List[Union[np.ndarray, torch.Tensor]], Union[np.ndarray, torch.Tensor]]
        ] = None,
        sampling_rate: Optional[int] = None,
        return_tensors: Optional[
            Union[str, transformers.TensorType]
        ] = transformers.TensorType.PYTORCH,
        include_audio_num_chunks: bool = False,
        images=None,
        images_kwargs: Optional[Dict[str, Any]] = None,
        **kwargs,
    ) -> transformers.BatchFeature:
        if text is not None and not isinstance(text, str):
            raise ValueError("Text must be a string. Batch mode not supported yet.")
        if audio is not None and audios is not None:
            raise ValueError("Only one of `audio` or `audios` should be provided.")
        if audio is not None:
            audios = audio
        if audios is None:
            audios = []
        elif isinstance(audios, (np.ndarray, torch.Tensor)):
            audios = [audios] if audios.ndim == 1 else list(audios)

        if text is not None and text.count(self.audio_placeholder) != len(audios):
            raise ValueError(
                f"Expected {len(audios)} audio placeholders, "
                f"found {text.count(self.audio_placeholder)}."
            )

        image_kwargs = dict(images_kwargs or {})
        for key in self.image_processor.valid_kwargs.__annotations__:
            if key in kwargs:
                image_kwargs[key] = kwargs.pop(key)
        nested_text_kwargs = kwargs.pop("text_kwargs", {}) or {}
        nested_audio_kwargs = kwargs.pop("audio_kwargs", {}) or {}
        text_kwargs = {**kwargs, **nested_text_kwargs}
        audio_kwargs = {**kwargs, **nested_audio_kwargs}

        data = {}
        image_token_counts = []
        has_images = images is not None and not (
            isinstance(images, (list, tuple)) and len(images) == 0
        )
        if has_images:
            if any(
                token not in self.vocab
                for token in (self.image_token, self.boi_token, self.eoi_token)
            ):
                raise ValueError("The tokenizer does not define Gemma4 image tokens.")
            image_kwargs["return_tensors"] = "pt"
            image_inputs = self.image_processor(images=images, **image_kwargs)
            image_token_counts = image_inputs.pop("num_soft_tokens_per_image").tolist()
            data.update(image_inputs)

        if text is not None:
            num_placeholders = text.count(self.image_token)
            if num_placeholders != len(image_token_counts):
                raise ValueError(
                    f"Expected {len(image_token_counts)} image placeholders, "
                    f"found {num_placeholders}."
                )
            replacements = iter(
                self.boi_token + self.image_token * count + self.eoi_token
                for count in image_token_counts
            )
            # Expand images before tokenizing audio spans so their offsets stay valid.
            text = re.sub(re.escape(self.image_token), lambda _: next(replacements), text)

        audio_is_continuation = []
        if len(audios) > 0:
            audios = [
                x.detach().cpu().numpy() if isinstance(x, torch.Tensor) else x
                for x in audios
            ]
            sr = sampling_rate or self.audio_processor.sampling_rate
            audios = _prepend_silence_to_short_audios(audios, sr)
            if text is not None and not include_audio_num_chunks:
                text, audios = self._auto_chunk_audios(text, audios, sr)
            hop_length = self.audio_processor.hop_length
            audios = [
                np.pad(x, (0, 2 * hop_length - len(x)), mode="constant")
                if len(x) < 2 * hop_length
                else x
                for x in audios
            ]
            audio_kwargs.pop("add_special_tokens", None)
            audio_kwargs.update(
                sampling_rate=sr,
                padding=self.audio_padding,
                pad_to_multiple_of=hop_length,
                truncation=False,
                return_attention_mask=True,
            )
            x: transformers.BatchFeature = self.audio_processor(audios, **audio_kwargs)
            data.update(
                self._chunk_and_pad_audio(
                    audio_values=torch.as_tensor(
                        x.input_features if "input_features" in x else x.input_values
                    ),
                    audio_lens=torch.as_tensor(x.attention_mask).sum(-1),
                    include_audio_num_chunks=include_audio_num_chunks,
                )
            )
            audio_is_continuation = data.pop("audio_is_continuation")
            data["audio_token_len"] = torch.ceil(
                data["audio_lens"] / (self.encoder_ds_factor * self.stack_factor)
            ).to(dtype=torch.int)

        if text is not None:
            text_kwargs["add_special_tokens"] = False
            tokenized_parts = self.tokenizer(
                text.split(self.audio_placeholder),
                **text_kwargs,
            )

            audio_token_start_idx = []
            placeholder_index = -1
            split_input_ids = tokenized_parts["input_ids"]
            input_ids: List[int] = []
            audio_replacement_token_id = self.vocab[self.audio_token_replacement]

            for i, token_len in enumerate(data.get("audio_token_len", [])):
                if not audio_is_continuation[i]:
                    placeholder_index += 1
                    if placeholder_index >= len(split_input_ids):
                        raise ValueError(f"Too few audio placeholders (expected {len(audios)}).")
                    input_ids.extend(split_input_ids[placeholder_index])
                audio_token_start_idx.append(len(input_ids))
                input_ids.extend([audio_replacement_token_id] * token_len)

            placeholder_index += 1
            if placeholder_index != len(split_input_ids) - 1:
                raise ValueError(f"Too many audio placeholders (expected {len(audios)}).")
            input_ids.extend(split_input_ids[placeholder_index])

            if "audio_token_len" in data:
                data["audio_token_start_idx"] = torch.as_tensor(audio_token_start_idx)

            data["input_ids"] = [input_ids]
            data["attention_mask"] = [[1] * len(input_ids)]
            if image_token_counts:
                if input_ids.count(self.image_token_id) != sum(image_token_counts):
                    raise ValueError("Tokenization changed the number of image placeholders.")
                mm_token_type_ids = [
                    int(token == self.image_token_id) for token in input_ids
                ]
                for start, length in zip(
                    audio_token_start_idx, data.get("audio_token_len", [])
                ):
                    mm_token_type_ids[start : start + int(length)] = [3] * int(length)
                data["mm_token_type_ids"] = [mm_token_type_ids]

        return transformers.BatchFeature(data=data, tensor_type=return_tensors)

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

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

    @property
    def model_input_names(self):
        return list(dict.fromkeys(self.tokenizer.model_input_names + [
            "audio_values", "audio_lens", "audio_token_len", "audio_token_start_idx",
            "audio_batch_size", "pixel_values", "image_position_ids", "mm_token_type_ids",
        ]))


AprielVoiceProcessor.register_for_auto_class()
transformers.AutoProcessor.register(AprielVoiceConfig, AprielVoiceProcessor)