from ...model.audio import BaseAudioModel
from ...model.engine import DiffusionPipeline, Engine
from ...model.vendor import TextGenerationVendor

from importlib import import_module
from typing import Any, Literal, cast

from torch import argmax, inference_mode
from transformers import (
    AutoModelForCTC,
    PreTrainedModel,
)


def _auto_processor() -> Any:
    return getattr(import_module("transformers"), "AutoProcessor")


class SpeechRecognitionModel(BaseAudioModel):
    _processor: Any

    def _load_model(
        self,
    ) -> PreTrainedModel | TextGenerationVendor | DiffusionPipeline:
        self._processor = cast(
            Any,
            cast(Any, _auto_processor()).from_pretrained(
                self._model_id,
                trust_remote_code=self._settings.trust_remote_code,
                # default behavior in transformers v4.48
                use_fast=True,
                subfolder=self._settings.tokenizer_subfolder or "",
            ),
        )
        assert self._model_id
        model = cast(
            PreTrainedModel,
            cast(Any, AutoModelForCTC).from_pretrained(
                self._model_id,
                trust_remote_code=self._settings.trust_remote_code,
                pad_token_id=self._processor.tokenizer.pad_token_id,
                ctc_loss_reduction="mean",
                device_map=self._device,
                tp_plan=Engine._get_tp_plan(self._settings.parallel),
                distributed_config=Engine._get_distributed_config(
                    self._settings.distributed_config
                ),
                ignore_mismatched_sizes=True,
                subfolder=self._settings.subfolder or "",
            ),
        )
        return model

    async def __call__(
        self,
        path: str,
        sampling_rate: int = 16_000,
        tensor_format: Literal["pt"] = "pt",
    ) -> str:
        audio = self._resample(path, sampling_rate)
        inputs = self._processor(
            audio,
            sampling_rate=sampling_rate,
            return_tensors=tensor_format,
        ).to(self._device)
        with inference_mode():
            # shape (batch, time_steps, vocab_size)
            logits = self._model(inputs.input_values).logits
        predicted_ids = argmax(logits, dim=-1)
        transcription = self._processor.batch_decode(predicted_ids)[0]
        return cast(str, transcription)
