"""Движки распознавания речи: whisper-large-v3-russian и Qwen3-ASR-1.7B."""
import warnings
import logging

import librosa
import numpy as np
import soundfile as sf
import torch
import torchaudio
from transformers import WhisperForConditionalGeneration, WhisperProcessor

from qwen_asr import Qwen3ASRModel

from .config import settings

logger = logging.getLogger(__name__)


class WhisperModel:
    """whisper-large-v3-russian (transformers)."""

    def __init__(self, device):
        self.model = WhisperForConditionalGeneration.from_pretrained(
            settings.WHISPER_MODEL_PATH
        )
        self.processor = WhisperProcessor.from_pretrained(settings.WHISPER_MODEL_PATH)
        self.device = device
        self.model.to(device)

    def transcribe(self, audio_path):
        wav, sr = sf.read(audio_path)
        wav = wav.mean(axis=1) if wav.ndim > 1 else wav
        wav16 = torchaudio.functional.resample(
            torch.from_numpy(np.asarray(wav, dtype=np.float32)), sr, 16000
        )
        inputs = self.processor(
            wav16, sampling_rate=16000, return_tensors="pt"
        )
        with torch.no_grad():
            ids = self.model.generate(inputs["input_features"].to(self.device))
        return self.processor.batch_decode(ids.cpu().numpy(), skip_special_tokens=True)


class QwenModel:
    """Qwen3-ASR-1.7B (qwen-asr)."""

    def __init__(self, device):
        self.model = Qwen3ASRModel.from_pretrained(settings.QWEN_MODEL_NAME)
        self.device = device

    def transcribe(self, audio_path):
        result = self.model.transcribe(audio_path)
        return result.get("text") if isinstance(result, dict) else result


def load_models(device):
    """Загрузить все активные модели."""
    models = {}
    if settings.LOAD_WHISPER:
        logger.info("Загрузка whisper-large-v3-russian...")
        models["whisper"] = WhisperModel(device)
    if settings.LOAD_QWEN:
        logger.info("Загрузка Qwen3-ASR-1.7B...")
        models["qwen3"] = QwenModel(device)
    return models