import io
import numpy as np
import torch
from pydub import AudioSegment
from transformers import (
    pipeline,
    WhisperProcessor,
    WhisperForConditionalGeneration,
)
from typing import Tuple

from app.config import config


class WhisperASR:
    """Explicit Whisper loader — avoids pipeline preprocessor num_frames bug."""

    def __init__(self, model_id: str, token: str, device: torch.device):
        self.processor = WhisperProcessor.from_pretrained(
            model_id,
            token=token,
        )
        self.model = WhisperForConditionalGeneration.from_pretrained(
            model_id,
            token=token,
            use_safetensors=False,
        ).to(device)
        self.model.eval()
        self.device = device

    def transcribe(self, samples: np.ndarray, sampling_rate: int) -> str:
        inputs = self.processor(
            samples,
            sampling_rate=sampling_rate,
            return_tensors="pt",
        ).to(self.device)

        with torch.no_grad():
            predicted_ids = self.model.generate(**inputs)

        transcription = self.processor.batch_decode(
            predicted_ids,
            skip_special_tokens=True,
        )
        return transcription[0].strip()


class STTPipeline:
    def __init__(self):
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        token = config.HF_TOKEN or None

        # Language identifier — wav2vec2, pipeline is fine here
        self.lid = pipeline(
            "audio-classification",
            model=config.LID_MODEL,
            device=0 if self.device.type == "cuda" else -1,
            token=token,
        )

        # ASR models — explicit Whisper loader per language
        self.asr: dict = {
            lang: WhisperASR(model_id, token, self.device)
            for lang, model_id in config.ASR_MODELS.items()
        }

    def decode_audio(self, audio_bytes: bytes) -> Tuple[np.ndarray, float]:
        seg = AudioSegment.from_file(io.BytesIO(audio_bytes))
        seg = seg.set_channels(1).set_frame_rate(config.SAMPLING_RATE)
        samples = np.array(seg.get_array_of_samples()).astype(np.float32)
        samples /= np.iinfo(seg.array_type).max
        duration = len(seg) / 1000.0
        return samples, duration

    def classify_language(self, samples: np.ndarray) -> Tuple[str, float]:
        result = self.lid(
            {"array": samples, "sampling_rate": config.SAMPLING_RATE},
            top_k=1,
        )
        label      = result[0]["label"].lower()
        confidence = result[0]["score"]
        return label, confidence

    def transcribe(self, audio_bytes: bytes) -> dict:
        samples, duration = self.decode_audio(audio_bytes)
        language, confidence = self.classify_language(samples)

        asr_model = self.asr.get(language)
        if asr_model is None:
            raise ValueError(
                f"No ASR model available for detected language: '{language}'. "
                f"Supported languages: {list(self.asr.keys())}"
            )

        transcription = asr_model.transcribe(samples, config.SAMPLING_RATE)
        return {
            "transcription": transcription,
            "language":      language,
            "confidence":    round(confidence, 4),
            "duration_sec":  round(duration, 2),
        }


stt_pipeline = STTPipeline()