"""Speaker verification: the assistant only answers Ahmed.

Engine: sherpa-onnx + NVIDIA TitaNet-small (192-dim embeddings, ~14ms
per 3s utterance on the M4; statically-linked onnxruntime, no torch).
Chosen after benchmarking 7 models on this machine — best speed/margin.

Lifecycle:
  ENROLLING  no profile yet: the first ENROLL_N good utterances (>= MIN_S
             net audio) build the profile. The assistant answers normally
             during enrollment (it's presumed to be the owner talking).
  LOCKED     each utterance is scored (cosine vs profile). Accept at
             T_ACCEPT; short utterances use a lenient threshold instead
             of hard-failing (embeddings are unreliable under ~2s).
             CONTINUOUS LEARNING: utterances scoring >= T_ADAPT are
             folded into the profile with a small anchored-EMA step —
             anchored to the original enrollment so absorbed mistakes
             can't destroy the profile (verified experimentally: a
             plain-averaged impostor collapses the margin; an anchored
             EMA one costs ~0.08 and recovers).

Env: VOICE_LOCK=0 disable · SPEAKER_THRESHOLD (0.50) ·
     SPEAKER_PROFILE (models/speaker_profile.npz) · SPEAKER_RESET=1
"""

from __future__ import annotations

import os
from pathlib import Path

import numpy as np

_MODELS_DIR = Path(__file__).resolve().parent.parent / "models"
MODEL_PATH = _MODELS_DIR / "titanet_small.onnx"

ENROLL_N = 5          # utterances to build the initial profile
MIN_S = 2.0           # net seconds for a reliable embedding
T_ACCEPT = float(os.environ.get("SPEAKER_THRESHOLD", "0.50"))
T_ACCEPT_SHORT = 0.35  # lenient gate for sub-2s utterances
T_ADAPT = 0.65        # only learn from clearly-owner utterances
EMA_ALPHA = 0.1       # small adaptation steps
ANCHOR_LAMBDA = 0.1   # weight of original enrollment in every update


def _norm(v: np.ndarray) -> np.ndarray:
    return v / max(float(np.linalg.norm(v)), 1e-9)


class SpeakerGate:
    def __init__(self) -> None:
        import sherpa_onnx

        cfg = sherpa_onnx.SpeakerEmbeddingExtractorConfig(
            model=str(MODEL_PATH), num_threads=1, provider="cpu"
        )
        assert cfg.validate(), "speaker model missing/invalid"
        self._ex = sherpa_onnx.SpeakerEmbeddingExtractor(cfg)
        self.profile_path = Path(
            os.environ.get("SPEAKER_PROFILE",
                           str(_MODELS_DIR / "speaker_profile.npz"))
        )
        self._enroll_embs: list[np.ndarray] = []
        self._anchor: np.ndarray | None = None   # original enrollment mean
        self._profile: np.ndarray | None = None  # adapted profile
        if os.environ.get("SPEAKER_RESET") == "1":
            self.profile_path.unlink(missing_ok=True)
        if self.profile_path.is_file():
            data = np.load(self.profile_path)
            self._anchor = data["anchor"]
            self._profile = data["profile"]

    # ---------- state ----------

    @property
    def enrolled(self) -> bool:
        return self._profile is not None

    @property
    def enroll_progress(self) -> tuple[int, int]:
        return len(self._enroll_embs), ENROLL_N

    # ---------- core ----------

    def _embed(self, audio: np.ndarray) -> np.ndarray:
        s = self._ex.create_stream()
        s.accept_waveform(sample_rate=16000, waveform=audio)
        s.input_finished()
        emb = np.asarray(self._ex.compute(s), dtype=np.float32)
        return _norm(emb)  # raw sherpa embeddings are NOT L2-normalized

    def _save(self) -> None:
        np.savez(self.profile_path, anchor=self._anchor, profile=self._profile)

    def check(self, audio: np.ndarray, speech_s: float) -> tuple[bool, float, str]:
        """Gate one utterance. Returns (accept, score, state).

        state: 'enrolling' | 'enrolled' | 'ok' | 'rejected' | 'short-ok'
        """
        emb = self._embed(audio)

        if not self.enrolled:
            if speech_s >= MIN_S:
                self._enroll_embs.append(emb)
                if len(self._enroll_embs) >= ENROLL_N:
                    self._anchor = _norm(
                        np.mean(np.stack(self._enroll_embs), axis=0))
                    self._profile = self._anchor.copy()
                    self._save()
                    return True, 1.0, "enrolled"
            return True, 1.0, "enrolling"

        score = float(np.dot(emb, self._profile))
        if speech_s < MIN_S:
            return score >= T_ACCEPT_SHORT, score, (
                "short-ok" if score >= T_ACCEPT_SHORT else "rejected")
        if score < T_ACCEPT:
            return False, score, "rejected"
        # continuous learning: anchored EMA on confident matches only
        if score >= T_ADAPT:
            stepped = _norm((1 - EMA_ALPHA) * self._profile + EMA_ALPHA * emb)
            self._profile = _norm(
                ANCHOR_LAMBDA * self._anchor + (1 - ANCHOR_LAMBDA) * stepped)
            self._save()
        return True, score, "ok"