from pathlib import Path
from typing import Any


class ProviderDeEmocaoHuggingFace:
    """Classifica a emoção de uma fala localmente com Transformers."""

    def __init__(self, modelo: str = "superb/wav2vec2-base-superb-er") -> None:
        from transformers import AutoFeatureExtractor, AutoModelForAudioClassification
        self.modelo = modelo
        self.processador = AutoFeatureExtractor.from_pretrained(modelo, local_files_only=True)
        self.classificador = AutoModelForAudioClassification.from_pretrained(
            modelo, local_files_only=True
        )
        self.classificador.eval()

    def analisar(self, arquivo: str | Path, inicio: float, fim: float) -> dict[str, Any]:
        import soundfile as sf
        import torch
        audio, taxa = sf.read(str(arquivo), start=max(0, int(inicio * 16000)),
                              stop=max(0, int(fim * 16000)), dtype="float32")
        if getattr(audio, "ndim", 1) > 1:
            audio = audio.mean(axis=1)
        entradas = self.processador(audio, sampling_rate=taxa, return_tensors="pt")
        with torch.no_grad():
            logits = self.classificador(**entradas).logits
        probabilidades = torch.softmax(logits[0], dim=-1)
        indice = int(torch.argmax(probabilidades))
        rotulo = self.classificador.config.id2label[indice]
        voz_aparente = {"non-neutral-female": "feminina",
                        "non-neutral-male": "masculina"}.get(rotulo)
        return {"emocao": self.classificador.config.id2label[indice],
                "confianca": float(probabilidades[indice]),
                "voz_aparente": voz_aparente,
                "confianca_voz": float(probabilidades[indice]) if voz_aparente else None}