"""Lightweight online speaker diarization.

The OS only gives us a single *mixed* stream of all remote participants (via
the loopback device), so we can't get per-person audio channels. Instead, for
each finished utterance we compute a voice fingerprint (ECAPA-TDNN speaker
embedding) and cluster it online: each new utterance is matched to the most
similar known speaker, or starts a new one ("Person 1", "Person 2", ...).

This is best-effort: very short utterances, people talking over each other, or
very similar voices will reduce accuracy. It is intended to *separate* the
conversation, not to identify people by name.
"""

import os
import warnings

import numpy as np


class SpeakerDiarizer:
    def __init__(self, model_dir="models/ecapa", similarity_threshold=0.40,
                 max_speakers=10, min_seconds=0.7, device="cpu"):
        self.similarity_threshold = float(similarity_threshold)
        self.max_speakers = int(max_speakers)
        self.min_samples = int(min_seconds * 16000)
        self.available = False
        self._model = None

        # speaker registry: list of dicts {centroid: np.ndarray, count: int, label: str}
        self.speakers = []

        try:
            warnings.filterwarnings("ignore")
            from speechbrain.inference.speaker import EncoderClassifier
            self._model = EncoderClassifier.from_hparams(
                source=model_dir, savedir=model_dir, run_opts={"device": device}
            )
            self.available = True
            print("✅ Speaker diarization ready (ECAPA voice fingerprints)")
        except Exception as e:
            print(f"⚠️  Diarization unavailable, falling back to single 'Them': {e}")

    def _embed(self, audio_np):
        import torch
        with warnings.catch_warnings():
            warnings.simplefilter("ignore")
            wav = torch.from_numpy(audio_np.astype("float32")).unsqueeze(0)
            emb = self._model.encode_batch(wav).squeeze().detach().cpu().numpy()
        norm = np.linalg.norm(emb)
        return emb / norm if norm > 0 else emb

    def identify(self, audio_np):
        """Return a speaker label for this utterance's audio (16k mono float32)."""
        if not self.available:
            return "Them"
        # too short to fingerprint reliably -> attribute to most recent speaker
        if audio_np.shape[0] < self.min_samples:
            return self.speakers[-1]["label"] if self.speakers else "Person 1"

        try:
            emb = self._embed(audio_np)
        except Exception:
            return "Them"

        if not self.speakers:
            return self._add_speaker(emb)

        sims = [float(np.dot(emb, s["centroid"])) for s in self.speakers]
        best = int(np.argmax(sims))

        if sims[best] >= self.similarity_threshold:
            self._update_speaker(best, emb)
            return self.speakers[best]["label"]

        if len(self.speakers) < self.max_speakers:
            return self._add_speaker(emb)

        # registry full: attach to nearest existing speaker
        self._update_speaker(best, emb)
        return self.speakers[best]["label"]

    def _add_speaker(self, emb):
        label = f"Person {len(self.speakers) + 1}"
        self.speakers.append({"centroid": emb, "count": 1, "label": label})
        return label

    def _update_speaker(self, idx, emb):
        s = self.speakers[idx]
        # running mean of embeddings, renormalised to the unit sphere
        c = (s["centroid"] * s["count"] + emb) / (s["count"] + 1)
        norm = np.linalg.norm(c)
        s["centroid"] = c / norm if norm > 0 else c
        s["count"] += 1

    def reset(self):
        self.speakers = []