from typing import List, Optional, Union

import numpy as np
from transformers import AutoProcessor, AutoTokenizer, WhisperFeatureExtractor
from transformers.feature_extraction_utils import BatchFeature

AUDIO_TOKENS_PER_SEGMENT = 750  # must match the projector output length in modeling_yor_audio.py
AUDIO_SPECIAL_TOKENS = ["<|audio_bos|>", "<|audio_eos|>", "<|AUDIO|>"]
SAMPLE_RATE = 16_000
NUM_MEL_BINS = 128
HOP_LENGTH = 160
CHUNK_LENGTH = 30


class YorAudioProcessor:
    audio_token = "<|AUDIO|>"
    audio_bos_token = "<|audio_bos|>"
    audio_eos_token = "<|audio_eos|>"

    def __init__(
        self,
        afriqueqwen_path: str = "McGill-NLP/AfriqueQwen3.5-9B-50Langs-Instruct-v1",
        feature_extractor: Optional[WhisperFeatureExtractor] = None,
    ):
        self.feature_extractor = feature_extractor or WhisperFeatureExtractor(
            feature_size=NUM_MEL_BINS,
            sampling_rate=SAMPLE_RATE,
            hop_length=HOP_LENGTH,
            chunk_length=CHUNK_LENGTH,
            n_fft=400,
        )
        self.tokenizer = AutoTokenizer.from_pretrained(afriqueqwen_path)
        self._add_audio_tokens()

    def _add_audio_tokens(self):
        missing = [t for t in AUDIO_SPECIAL_TOKENS if t not in self.tokenizer.get_vocab()]
        if missing:
            self.tokenizer.add_special_tokens({"additional_special_tokens": missing})

    @property
    def audio_token_id(self) -> int:
        return self.tokenizer.convert_tokens_to_ids(self.audio_token)

    def build_prompt(self, text: str, num_audios: int = 1) -> str:
        audio_block = (
            self.audio_bos_token
            + self.audio_token * AUDIO_TOKENS_PER_SEGMENT
            + self.audio_eos_token
        )
        return audio_block * num_audios + text

    def __call__(
        self,
        text: Union[str, List[str]],
        audio: Optional[Union[np.ndarray, List[np.ndarray]]] = None,
        sampling_rate: int = SAMPLE_RATE,
        return_tensors: str = "pt",
        padding: bool = True,
        truncation: bool = True,
        max_length: Optional[int] = None,
    ) -> BatchFeature:
        if isinstance(text, str):
            text = [text]

        encoding = self.tokenizer(
            text,
            return_tensors=return_tensors,
            padding=padding,
            truncation=truncation,
            max_length=max_length,
        )
        result = BatchFeature(data=dict(encoding))

        if audio is not None:
            if isinstance(audio, np.ndarray):
                audio = [audio]
            if audio[0].ndim == 2:
                import torch
                result["input_features"] = torch.tensor(np.stack(audio))
            else:
                audio_out = self.feature_extractor(
                    audio, sampling_rate=sampling_rate, return_tensors=return_tensors
                )
                result["input_features"] = audio_out["input_features"]

        return result

    def decode(self, token_ids, **kwargs) -> str:
        return self.tokenizer.decode(token_ids, **kwargs)

    def batch_decode(self, sequences, **kwargs) -> List[str]:
        return self.tokenizer.batch_decode(sequences, **kwargs)

    def save_pretrained(self, save_directory: str):
        self.feature_extractor.save_pretrained(save_directory)
        self.tokenizer.save_pretrained(save_directory)

    @classmethod
    def from_pretrained(cls, path: str, **kwargs) -> "YorAudioProcessor":
        processor = cls.__new__(cls)
        processor.feature_extractor = WhisperFeatureExtractor.from_pretrained(path)
        processor.tokenizer = AutoTokenizer.from_pretrained(path)
        processor._add_audio_tokens()
        return processor

    @classmethod
    def register_for_auto_class(cls, auto_class="AutoProcessor"):
        pass


AutoProcessor.register("yor_audio", YorAudioProcessor)