#!/usr/bin/env python3
"""On-device speaker diarization + voiceprint identification for the local Diary path.

Engine: FluidAudio CoreML (council recapmuse-diarization-engine, 2026-07-31, 5/5), via the
`diarize-cli` helper binary (tools/diarize-cli). The offline community-1 pipeline scored ABOVE
the pyannote-3.1 python baseline on the spike harness (2spk 100.0% 8/8, 3spk 97.8% 11/12), and
its 256-dim wespeaker embeddings are the SAME space as the enrolled voiceprint library
(cross-stack cosine 0.98 for the same speaker, measured 2026-07-31).

Zero egress by construction at run time: the helper loads CoreML models from disk; audio and
embeddings never leave the machine. Model files are fetched ONCE (dev-time / setup-time), the
same pattern as the Parakeet ASR model.

  local_diarize.py <audio> [--num-speakers N]   JSON: {ok, turns:[{speaker,start,end}], names:{...}}
  local_diarize.py --check                       JSON: {ok, cli, models} availability probe
  local_diarize.py --enroll "Name" <audio> [--speaker SID]
                                                 enroll the dominant (or chosen) speaker's
                                                 voiceprint from a recording into the library

Env:
  EPIPHANI_DIARIZE_CLI     path to diarize-cli (default: bundled next to engine, then repo build)
  EPIPHANI_DIARIZER_MODELS local CoreML models dir (optional; default = FluidAudio cache)
  EPIPHANI_VOICEPRINTS     voiceprint library JSON (default: canonical library)
  EPIPHANI_VP_THRESHOLD    cosine threshold for naming a speaker (default 0.60)
"""
import json
import math
import os
import subprocess
import sys

_HERE = os.path.dirname(os.path.abspath(__file__))

VOICEPRINTS = os.path.expanduser(
    os.environ.get("EPIPHANI_VOICEPRINTS")
    or "~/.claude/artifacts/hvn-voiceprints/voiceprints.json")
# Calibration (2026-07-31): max BETWEEN-speaker cosine in the 6-person library is 0.442
# (two similar voices); same-speaker cross-stack identity measured 0.98 on clean audio.
# 0.60 names a speaker only with real margin; below it the generic label is kept.
VP_THRESHOLD = float(os.environ.get("EPIPHANI_VP_THRESHOLD") or 0.60)

# FluidAudio's offline models live in its own cache dir by default; the helper handles it.
_MODELS_DIR = os.environ.get("EPIPHANI_DIARIZER_MODELS") or ""


def _cli_path():
    p = os.environ.get("EPIPHANI_DIARIZE_CLI")
    if p and os.path.exists(os.path.expanduser(p)):
        return os.path.expanduser(p)
    candidates = [
        os.path.join(_HERE, "diarize-cli"),                                   # bundled in the app
        os.path.join(_HERE, "..", "..", "tools", "diarize-cli", ".build", "release", "diarize-cli"),  # repo dev build
    ]
    for c in candidates:
        c = os.path.abspath(c)
        if os.path.exists(c):
            return c
    return None


def _models_present():
    """True if the offline CoreML model set is already on disk (no download needed)."""
    if _MODELS_DIR:
        return os.path.isdir(os.path.expanduser(_MODELS_DIR))
    cache = os.path.expanduser("~/Library/Application Support/FluidAudio/Models")
    if not os.path.isdir(cache):
        return False
    for root, dirs, files in os.walk(cache):
        if any(f.endswith(".mlmodelc") or d.endswith(".mlmodelc") for f in files for d in [f]) or \
           any(d.endswith(".mlmodelc") for d in dirs):
            return True
    return False


def _cos(a, b):
    dot = sum(x * y for x, y in zip(a, b))
    na = math.sqrt(sum(x * x for x in a))
    nb = math.sqrt(sum(x * x for x in b))
    return dot / (na * nb) if na and nb else 0.0


def _load_library():
    try:
        with open(VOICEPRINTS) as fh:
            lib = json.load(fh)
        return {name: sp["embedding"] for name, sp in (lib.get("speakers") or {}).items()
                if isinstance(sp.get("embedding"), list)}
    except Exception:  # noqa: BLE001 — no library = generic labels, never a failure
        return {}


def identify(speakers, threshold=VP_THRESHOLD):
    """Map engine speaker ids to enrolled names. speakers = {id: {"embedding": [...], ...}}.
    Returns {id: name-or-None}. Greedy best-match, one enrolled name per recording."""
    lib = _load_library()
    out = {sid: None for sid in speakers}
    if not lib:
        return out
    scored = []
    for sid, info in speakers.items():
        emb = info.get("embedding") or []
        if len(emb) != len(next(iter(lib.values()))):
            continue                      # dimension mismatch: never guess
        for name, ref in lib.items():
            scored.append((_cos(emb, ref), sid, name))
    scored.sort(reverse=True)
    used_names, used_sids = set(), set()
    for score, sid, name in scored:
        if score < threshold:
            break
        if sid in used_sids or name in used_names:
            continue
        out[sid] = name
        used_sids.add(sid)
        used_names.add(name)
    return out


def diarize(audio_path, num_speakers=None):
    """Run the on-device diarizer. Returns {ok, turns, speakers, names} or {ok:False, error}."""
    cli = _cli_path()
    if not cli:
        return {"ok": False, "error": "diarize-cli not found (build tools/diarize-cli)"}
    args = [cli, audio_path]
    if num_speakers:
        args += ["--num-speakers", str(num_speakers)]
    env = {**os.environ}
    if _MODELS_DIR:
        env["EPIPHANI_DIARIZER_MODELS"] = _MODELS_DIR
    try:
        out = subprocess.run(args, capture_output=True, text=True, timeout=1800, env=env)
    except subprocess.TimeoutExpired:
        return {"ok": False, "error": "diarization timed out"}
    try:
        data = json.loads(out.stdout or "{}")
    except ValueError:
        return {"ok": False, "error": "diarizer returned no JSON: %s" % (out.stderr or "")[:200]}
    if not data.get("ok"):
        return {"ok": False, "error": data.get("error") or "diarization failed"}
    data["names"] = identify(data.get("speakers") or {})
    return data


def enroll(audio_path, name, speaker_id=None):
    """Enroll a voiceprint from a recording, fully on-device. Diarizes the audio, takes the
    chosen speaker (default: the one who talked longest), and folds its embedding into the
    library as a running average, so repeat enrollments strengthen the print. What is stored:
    ONLY a 256-number voice signature and a sample count, in the local library JSON file —
    never audio, never words. Returns {ok, name, samples_count, cosine_vs_previous}."""
    d = diarize(audio_path)
    if not d.get("ok"):
        return d
    speakers = d.get("speakers") or {}
    if not speakers:
        return {"ok": False, "error": "no speakers found in this recording"}
    sid = speaker_id or max(speakers, key=lambda s: speakers[s].get("duration") or 0)
    if sid not in speakers:
        return {"ok": False, "error": "speaker id %r not in %s" % (sid, sorted(speakers))}
    emb = speakers[sid].get("embedding") or []
    if not emb:
        return {"ok": False, "error": "chosen speaker has no embedding"}

    try:
        with open(VOICEPRINTS) as fh:
            lib = json.load(fh)
    except Exception:  # noqa: BLE001
        lib = {"version": 1, "embedder": "wespeaker-voxceleb-resnet34-LM (256d)", "speakers": {}}
    lib.setdefault("speakers", {})
    prev = lib["speakers"].get(name)
    cos_prev = None
    if prev and isinstance(prev.get("embedding"), list) and len(prev["embedding"]) == len(emb):
        cos_prev = round(_cos(emb, prev["embedding"]), 3)
        n = int(prev.get("samples_count") or 1)
        merged = [(p * n + e) / (n + 1) for p, e in zip(prev["embedding"], emb)]
        lib["speakers"][name] = {**prev, "embedding": merged, "samples_count": n + 1}
    else:
        lib["speakers"][name] = {"embedding": emb, "samples_count": 1, "reference_meetings": []}
    os.makedirs(os.path.dirname(VOICEPRINTS), exist_ok=True)
    with open(VOICEPRINTS, "w") as fh:
        json.dump(lib, fh)
    return {"ok": True, "name": name, "speaker_used": sid,
            "samples_count": lib["speakers"][name]["samples_count"],
            "cosine_vs_previous": cos_prev, "library": VOICEPRINTS}


def main(argv):
    if "--check" in argv:
        cli = _cli_path()
        sys.stdout.write(json.dumps({
            "ok": bool(cli) and _models_present(),
            "cli": cli, "models": _models_present(),
            "library_speakers": len(_load_library())}))
        return
    if "--enroll" in argv:
        name = argv[argv.index("--enroll") + 1]
        paths = [a for a in argv if not a.startswith("--") and a != name]
        sid = argv[argv.index("--speaker") + 1] if "--speaker" in argv else None
        if not paths:
            sys.stdout.write(json.dumps({"ok": False, "error": "usage: local_diarize.py --enroll \"Name\" <audio>"}))
            return
        sys.stdout.write(json.dumps(enroll(paths[0], name, speaker_id=sid)))
        return
    paths = [a for a in argv if not a.startswith("--")]
    if not paths:
        sys.stdout.write(json.dumps({"ok": False, "error": "usage: local_diarize.py <audio> | --check | --enroll"}))
        return
    n = None
    if "--num-speakers" in argv:
        n = int(argv[argv.index("--num-speakers") + 1])
    sys.stdout.write(json.dumps(diarize(paths[0], num_speakers=n)))


if __name__ == "__main__":
    main(sys.argv[1:])
