"""
Optional speaker diarization module for WhisperLive.

Uses speaker embeddings and online clustering to assign speaker labels
to transcription segments in real-time. Requires pyannote.audio as an
optional dependency.

Install: pip install pyannote.audio
"""

import logging
import numpy as np


def load_audio(file_path, sample_rate=16000):
    """Load an audio file as mono float32 PCM at the requested sample rate."""
    import av

    container = av.open(file_path)
    resampler = av.AudioResampler(format="flt", layout="mono", rate=sample_rate)
    chunks = []

    try:
        for frame in container.decode(audio=0):
            for resampled_frame in resampler.resample(frame):
                chunks.append(
                    resampled_frame.to_ndarray().reshape(-1).astype(np.float32)
                )
    finally:
        container.close()

    if not chunks:
        return np.array([], dtype=np.float32)
    return np.concatenate(chunks)


class SpeakerDiarizer:
    """Real-time speaker diarization using speaker embeddings and online clustering.

    Each completed transcription segment's audio is passed through a speaker
    embedding model. The embedding is compared against known speakers using
    cosine similarity. If no match exceeds the threshold, a new speaker is
    created.

    Args:
        similarity_threshold (float): Minimum cosine similarity to match an
            existing speaker. Lower values merge speakers more aggressively.
            Default 0.55.
        max_speakers (int): Maximum number of distinct speakers to track.
            Once reached, new segments are assigned to the closest existing
            speaker. Default 10.
        embedding_model (str): The pyannote embedding model to use.
            Default "pyannote/wespeaker-voxceleb-resnet34-LM".
        hf_token (str or None): HuggingFace token for gated model access.
    """

    def __init__(
        self,
        similarity_threshold=0.55,
        max_speakers=10,
        embedding_model="pyannote/wespeaker-voxceleb-resnet34-LM",
        hf_token=None,
        speaker_names=None,
    ):
        self.similarity_threshold = similarity_threshold
        self.max_speakers = max_speakers
        self.speaker_names = list(speaker_names or [])
        self.speakers = {}  # speaker_id -> embedding (averaged)
        self._speaker_count = 0
        self._model = None
        self._embedding_model_name = embedding_model
        self._hf_token = hf_token

    def _next_speaker_id(self):
        if self._speaker_count < len(self.speaker_names):
            return self.speaker_names[self._speaker_count]
        return f"SPEAKER_{self._speaker_count:02d}"

    def _load_model(self):
        """Lazy-load the embedding model on first use."""
        if self._model is not None:
            return
        try:
            from pyannote.audio import Model, Inference
            import torch

            model = Model.from_pretrained(
                self._embedding_model_name,
                use_auth_token=self._hf_token,
            )
            device = "cuda" if torch.cuda.is_available() else "cpu"
            self._model = Inference(model, window="whole", device=torch.device(device))
            logging.info(f"Speaker embedding model loaded on {device}")
        except ImportError:
            raise ImportError(
                "pyannote.audio is required for speaker diarization. "
                "Install it with: pip install pyannote.audio"
            )

    def _compute_embedding(self, audio_np, sample_rate=16000):
        """Compute a speaker embedding from an audio numpy array.

        Args:
            audio_np (np.ndarray): 1-D float32 audio samples.
            sample_rate (int): Sample rate of the audio.

        Returns:
            np.ndarray: Speaker embedding vector, or None if audio is too short.
        """
        self._load_model()
        if len(audio_np) < sample_rate * 0.3:
            return None
        waveform = {
            "waveform": __import__("torch").tensor(audio_np).unsqueeze(0),
            "sample_rate": sample_rate,
        }
        embedding = self._model(waveform)
        return embedding / np.linalg.norm(embedding)

    @staticmethod
    def _cosine_similarity(a, b):
        """Compute cosine similarity between two vectors."""
        return float(np.dot(a, b))

    def identify_speaker(self, audio_np, sample_rate=16000):
        """Identify or create a speaker from an audio segment.

        Args:
            audio_np (np.ndarray): 1-D float32 audio for the segment.
            sample_rate (int): Sample rate. Default 16000.

        Returns:
            str or None: Speaker label (e.g. "SPEAKER_00"), or None if
                the audio is too short to embed.
        """
        embedding = self._compute_embedding(audio_np, sample_rate)
        if embedding is None:
            return None

        best_speaker = None
        best_sim = -1.0

        for speaker_id, stored_emb in self.speakers.items():
            sim = self._cosine_similarity(embedding, stored_emb)
            if sim > best_sim:
                best_sim = sim
                best_speaker = speaker_id

        if best_sim >= self.similarity_threshold:
            # Update running average for the matched speaker
            self.speakers[best_speaker] = (
                self.speakers[best_speaker] * 0.9 + embedding * 0.1
            )
            # Re-normalize
            self.speakers[best_speaker] /= np.linalg.norm(self.speakers[best_speaker])
            return best_speaker

        if len(self.speakers) >= self.max_speakers:
            # Assign to closest speaker
            return (
                best_speaker if best_speaker else f"SPEAKER_{self._speaker_count:02d}"
            )

        # Create a new speaker
        speaker_id = self._next_speaker_id()
        self._speaker_count += 1
        self.speakers[speaker_id] = embedding
        return speaker_id

    def enroll_speaker(self, speaker_name, audio_np, sample_rate=16000):
        """Enroll a known speaker from reference audio."""
        embedding = self._compute_embedding(audio_np, sample_rate)
        if embedding is None:
            return False
        self.speakers[speaker_name] = embedding
        return True

    def reset(self):
        """Reset all speaker state."""
        self.speakers.clear()
        self._speaker_count = 0