"""
SpeechBrain + librosa vocal analysis.
Performs emotion classification and acoustic feature extraction from audio.
"""
import io
import logging
import numpy as np
import soundfile as sf
import librosa
from speechbrain.inference.interfaces import foreign_class
from dataclasses import dataclass, asdict
from typing import Optional

logger = logging.getLogger(__name__)

SAMPLE_RATE = 16000
MIN_DURATION_SECONDS = 0.5


@dataclass
class VocalAnalysis:
    emotion: str
    emotion_confidence: float
    emotion_scores: dict
    pitch_mean_hz: Optional[float]
    pitch_variance: Optional[float]
    energy_rms: float
    speaking_rate_approx: float
    voiced_fraction: float
    duration_seconds: float
    analysis_error: Optional[str] = None


class VocalAnalyzer:
    """Loads models once at startup and provides analysis per audio chunk."""

    def __init__(self, model_dir: str = "/app/models"):
        logger.info("Loading SpeechBrain emotion model...")
        self.classifier = foreign_class(
            source="speechbrain/emotion-recognition-wav2vec2-IEMOCAP",
            pymodule_file="custom_interface.py",
            classname="CustomEncoderWav2vec2Classifier",
            savedir=model_dir,
            run_opts={"device": "cpu"},
        )
        logger.info("SpeechBrain model loaded successfully.")

    def analyze(self, audio_bytes: bytes, content_type: str = "audio/wav") -> VocalAnalysis:
        try:
            audio_buf = io.BytesIO(audio_bytes)
            audio, sr = sf.read(audio_buf, dtype="float32")

            if audio.ndim > 1:
                audio = audio.mean(axis=1)

            if sr != SAMPLE_RATE:
                audio = librosa.resample(audio, orig_sr=sr, target_sr=SAMPLE_RATE)
                sr = SAMPLE_RATE

            duration = len(audio) / sr

            if duration < MIN_DURATION_SECONDS:
                return VocalAnalysis(
                    emotion="neutral", emotion_confidence=0.0, emotion_scores={},
                    pitch_mean_hz=None, pitch_variance=None, energy_rms=0.0,
                    speaking_rate_approx=0.0, voiced_fraction=0.0,
                    duration_seconds=duration,
                    analysis_error=f"Audio too short: {duration:.2f}s < {MIN_DURATION_SECONDS}s minimum",
                )

            import tempfile, os
            with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as tmp:
                sf.write(tmp.name, audio, sr)
                tmp_path = tmp.name

            try:
                out_prob, score, index, text_lab = self.classifier.classify_file(tmp_path)
                emotion = text_lab[0] if isinstance(text_lab, list) else text_lab
                confidence = float(score[0]) if hasattr(score, "__iter__") else float(score)
                labels = ["neu", "hap", "sad", "ang", "fru"]
                prob_list = out_prob[0].tolist() if hasattr(out_prob[0], "tolist") else list(out_prob[0])
                emotion_scores = {l: round(p, 4) for l, p in zip(labels, prob_list)}
            finally:
                os.unlink(tmp_path)

            f0, voiced_flag, _ = librosa.pyin(
                audio, fmin=librosa.note_to_hz("C2"), fmax=librosa.note_to_hz("C7"), sr=sr,
            )

            voiced_f0 = f0[voiced_flag] if voiced_flag is not None else np.array([])
            pitch_mean = float(np.mean(voiced_f0)) if len(voiced_f0) > 0 else None
            pitch_var = float(np.var(voiced_f0)) if len(voiced_f0) > 1 else None
            voiced_fraction = float(np.sum(voiced_flag) / len(voiced_flag)) if voiced_flag is not None else 0.0

            rms = librosa.feature.rms(y=audio)
            energy_rms = float(np.mean(rms))

            onset_frames = librosa.onset.onset_detect(y=audio, sr=sr)
            speaking_rate = len(onset_frames) / duration if duration > 0 else 0.0

            return VocalAnalysis(
                emotion=emotion, emotion_confidence=round(confidence, 4),
                emotion_scores=emotion_scores,
                pitch_mean_hz=round(pitch_mean, 2) if pitch_mean else None,
                pitch_variance=round(pitch_var, 2) if pitch_var else None,
                energy_rms=round(energy_rms, 4),
                speaking_rate_approx=round(speaking_rate, 2),
                voiced_fraction=round(voiced_fraction, 4),
                duration_seconds=round(duration, 3),
            )

        except Exception as e:
            logger.exception(f"Analysis failed: {e}")
            return VocalAnalysis(
                emotion="neutral", emotion_confidence=0.0, emotion_scores={},
                pitch_mean_hz=None, pitch_variance=None, energy_rms=0.0,
                speaking_rate_approx=0.0, voiced_fraction=0.0,
                duration_seconds=0.0, analysis_error=str(e),
            )