from __future__ import annotations

import json
import math
from dataclasses import asdict, dataclass, field
from pathlib import Path
from time import time
from typing import Iterable


def cosine_similarity(left: Iterable[float], right: Iterable[float]) -> float:
    a = [float(x) for x in left]
    b = [float(x) for x in right]
    if len(a) != len(b) or not a:
        return 0.0
    dot = sum(x * y for x, y in zip(a, b, strict=True))
    norm_a = math.sqrt(sum(x * x for x in a))
    norm_b = math.sqrt(sum(y * y for y in b))
    if norm_a == 0.0 or norm_b == 0.0:
        return 0.0
    return dot / (norm_a * norm_b)


@dataclass(slots=True)
class SpeakerEmbedding:
    vector: list[float]
    source: str
    created_at: float = field(default_factory=time)


@dataclass(slots=True)
class SpeakerProfile:
    speaker_id: str
    display_name: str
    embeddings: list[SpeakerEmbedding] = field(default_factory=list)


@dataclass(frozen=True, slots=True)
class SpeakerMatch:
    speaker_id: str | None
    confidence: float
    display_name: str | None = None


class SpeakerStore:
    """Small JSON-backed speaker profile store for the MVP.

    This intentionally stores embeddings, not raw audio. Raw clip retention can be
    added later for unknown-speaker review.
    """

    def __init__(self, path: Path):
        self.path = path
        self.profiles: dict[str, SpeakerProfile] = {}
        self._load()

    def add_embedding(
        self,
        speaker_id: str,
        vector: Iterable[float],
        source: str,
        display_name: str | None = None,
    ) -> None:
        profile = self.profiles.get(speaker_id)
        if profile is None:
            profile = SpeakerProfile(
                speaker_id=speaker_id,
                display_name=display_name or speaker_id,
            )
            self.profiles[speaker_id] = profile
        profile.embeddings.append(
            SpeakerEmbedding(vector=[float(x) for x in vector], source=source)
        )
        self.save()

    def identify(self, vector: Iterable[float]) -> SpeakerMatch:
        probe = [float(x) for x in vector]
        best = SpeakerMatch(speaker_id=None, confidence=0.0)
        for profile in self.profiles.values():
            for embedding in profile.embeddings:
                score = cosine_similarity(probe, embedding.vector)
                if score > best.confidence:
                    best = SpeakerMatch(
                        speaker_id=profile.speaker_id,
                        display_name=profile.display_name,
                        confidence=score,
                    )
        return best

    def save(self) -> None:
        self.path.parent.mkdir(parents=True, exist_ok=True)
        data = {speaker_id: asdict(profile) for speaker_id, profile in self.profiles.items()}
        self.path.write_text(json.dumps(data, indent=2), encoding="utf-8")

    def _load(self) -> None:
        if not self.path.exists():
            return
        raw = json.loads(self.path.read_text(encoding="utf-8"))
        for speaker_id, profile_data in raw.items():
            embeddings = [
                SpeakerEmbedding(**embedding)
                for embedding in profile_data.get("embeddings", [])
            ]
            self.profiles[speaker_id] = SpeakerProfile(
                speaker_id=profile_data.get("speaker_id", speaker_id),
                display_name=profile_data.get("display_name", speaker_id),
                embeddings=embeddings,
            )