#!/usr/bin/env python3
"""On-device transcription for Epiphani's local organizer path (NO network, NO cloud).

This is the on-device counterpart to organize.py's cloud `_transcribe` (which posts audio to
OpenAI). It transcribes one or more recordings entirely on this Mac using Parakeet (MLX) with a
model that is already cached locally, so with the network disabled it still produces a transcript.

Runs under the pyenv Python that has `parakeet_mlx` + `mlx` installed (NOT the system python that
runs organize.py). organize.py shells out to this script, same pattern models_data.add() uses for
huggingface downloads. The model is loaded ONCE per invocation, so pass every clip in one call.

Integrity: HF_HUB_OFFLINE is forced on so a cached model can never trigger a silent download; if the
model is not cached or the LLM/ASR stack is missing, this fails visibly (ok:false) and organize.py
must surface that rather than fall back to the cloud.

  local_transcribe.py <audio1> [<audio2> ...]     # JSON: {ok, model, results:{path:{ok,transcript}}}
  local_transcribe.py --diarize <audio1> [...]     # adds on-device speaker labels + names per clip
  local_transcribe.py --check                      # JSON: {ok, model, cached} availability probe

Diarization (--diarize) runs the FluidAudio CoreML pipeline via local_diarize.py (council
recapmuse-diarization-engine 2026-07-31), merges its speaker turns with Parakeet's sentence
timestamps, and labels lines "Name:" (enrolled voiceprint match) or "Speaker A:" (unknown),
matching the cloud path's transcript shape. Fully on-device; results add {speakers: [...]}.

Env:
  EPIPHANI_ASR_MODEL   Parakeet repo id (default mlx-community/parakeet-tdt-0.6b-v3)
  EPIPHANI_CAP_SEC     transcribe at most the first N seconds/clip (default 600), parity with cloud
"""
import json
import os
import subprocess
import sys

# Force fully offline: a cached model is used; a missing one fails loudly, never downloads.
os.environ.setdefault("HF_HUB_OFFLINE", "1")
os.environ.setdefault("TRANSFORMERS_OFFLINE", "1")

ASR_MODEL = os.environ.get("EPIPHANI_ASR_MODEL") or "mlx-community/parakeet-tdt-0.6b-v3"
CAP_SEC = int(os.environ.get("EPIPHANI_CAP_SEC") or 600)


def _ffmpeg():
    for p in ("/opt/homebrew/bin/ffmpeg", "/usr/local/bin/ffmpeg", "/usr/bin/ffmpeg"):
        if os.path.exists(p):
            return p
    from shutil import which
    return which("ffmpeg")


FF = _ffmpeg()


def _model_cached():
    """True if the Parakeet model is already in the HF hub cache (so offline transcription works)."""
    hub = os.path.expanduser("~/.cache/huggingface/hub")
    tag = "models--" + ASR_MODEL.replace("/", "--")
    return os.path.isdir(os.path.join(hub, tag))


def _prep(path, idx):
    """Decode+trim a recording to a clean 16 kHz mono wav parakeet likes; cap at CAP_SEC (cloud parity)."""
    tmp = "/tmp/epiphani_local_%d_%d.wav" % (os.getpid(), idx)
    args = [FF, "-y", "-hide_banner", "-loglevel", "error", "-nostdin", "-i", path,
            "-t", str(CAP_SEC), "-ac", "1", "-ar", "16000", "-f", "wav", tmp]
    subprocess.run(args, check=True)
    return tmp


def _overlap(a0, a1, b0, b1):
    return max(0.0, min(a1, b1) - max(a0, b0))


def _diarized_transcript(sentences, turns, names):
    """Merge Parakeet sentences with diarization turns into the cloud path's transcript shape.

    sentences: [{start, end, text}]; turns: [{speaker, start, end}]; names: {engine_id: name|None}.
    Each sentence gets the speaker with the largest time overlap (nearest turn as fallback).
    Returns (transcript_lines, ordered_speaker_labels)."""
    display = {}
    letters = iter("ABCDEFGHIJKLMNOPQRSTUVWXYZ")

    def label_for(sid):
        if sid not in display:
            display[sid] = names.get(sid) or ("Speaker %s" % next(letters, "?"))
        return display[sid]

    lines, speakers, cur, buf = [], [], None, []
    for s in sentences:
        best, best_ov = None, 0.0
        for t in turns:
            ov = _overlap(s["start"], s["end"], t["start"], t["end"])
            if ov > best_ov:
                best, best_ov = t["speaker"], ov
        if best is None and turns:      # no overlap at all: nearest turn midpoint
            mid = (s["start"] + s["end"]) / 2
            best = min(turns, key=lambda t: min(abs(mid - t["start"]), abs(mid - t["end"])))["speaker"]
        lab = label_for(best) if best is not None else "Speaker ?"
        if lab not in speakers:
            speakers.append(lab)
        if lab != cur and buf:
            lines.append("%s: %s" % (cur, " ".join(buf)))
            buf = []
        cur = lab
        txt = (s.get("text") or "").strip()
        if txt:
            buf.append(txt)
    if buf:
        lines.append("%s: %s" % (cur, " ".join(buf)))
    labeled_turns = [{"speaker": label_for(t["speaker"]), "start": t["start"], "end": t["end"]}
                     for t in turns]
    return "\n".join(lines), speakers, labeled_turns


def transcribe_all(paths, diarize=False):
    if not FF:
        return {"ok": False, "error": "ffmpeg not found (brew install ffmpeg)"}
    if not _model_cached():
        return {"ok": False,
                "error": "on-device model not downloaded: %s. Add it in Settings > AI Engines." % ASR_MODEL}
    try:
        from parakeet_mlx import from_pretrained
    except Exception as e:  # noqa: BLE001
        return {"ok": False, "error": "parakeet_mlx not installed in this Python: %s" % e}
    try:
        model = from_pretrained(ASR_MODEL)          # loaded once for the whole batch
    except Exception as e:  # noqa: BLE001
        return {"ok": False, "error": "could not load on-device model %s: %s" % (ASR_MODEL, str(e)[:200])}

    results = {}
    for i, path in enumerate(paths):
        tmp = None
        try:
            tmp = _prep(path, i)
            r = model.transcribe(tmp, chunk_duration=300.0)   # chunk long clips so memory stays bounded
            plain = (getattr(r, "text", "") or "").strip()
            entry = {"ok": True, "transcript": plain}
            if diarize and plain:
                # Same prepped wav feeds both models, so timestamps line up exactly.
                import local_diarize
                d = local_diarize.diarize(tmp)
                if d.get("ok") and d.get("turns"):
                    sentences = [{"start": float(s.start), "end": float(s.end), "text": s.text}
                                 for s in (getattr(r, "sentences", None) or [])]
                    if sentences:
                        transcript, speakers, labeled_turns = _diarized_transcript(
                            sentences, d["turns"], d.get("names") or {})
                        if transcript.strip():
                            entry.update(transcript=transcript, speakers=speakers,
                                         turns=labeled_turns, diarized=True)
                else:
                    entry["diarize_error"] = d.get("error") or "diarization unavailable"
            results[path] = entry
        except Exception as e:  # noqa: BLE001
            results[path] = {"ok": False, "error": str(e)[:200]}
        finally:
            if tmp and os.path.exists(tmp):
                try:
                    os.remove(tmp)
                except OSError:
                    pass
    return {"ok": True, "model": ASR_MODEL, "results": results}


def main(argv):
    if "--check" in argv:
        cached = _model_cached()
        ok = cached
        try:
            import parakeet_mlx  # noqa: F401
            have_pkg = True
        except Exception:  # noqa: BLE001
            have_pkg = False
        sys.stdout.write(json.dumps({
            "ok": bool(ok and have_pkg and FF), "model": ASR_MODEL, "cached": cached,
            "parakeet": have_pkg, "ffmpeg": bool(FF)}))
        return
    paths = [a for a in argv if not a.startswith("--")]
    if not paths:
        sys.stdout.write(json.dumps({"ok": False, "error": "usage: local_transcribe.py [--diarize] <audio> [...] | --check"}))
        return
    sys.stdout.write(json.dumps(transcribe_all(paths, diarize="--diarize" in argv)))


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