#!/usr/bin/env python3
"""Run pyannote speaker diarization on one audio file and write the turns as JSON.

Usage: diarize.py AUDIO OUT_JSON [--speakers N | --min-speakers N --max-speakers N]

Output is a list of {"start", "end", "speaker"} in seconds, which is what
merge_transcript.py reads. HF_TOKEN is only needed the first time, to download
the gated model; after that the cached copy loads offline.

pyannote and torch are imported inside run(), so the pure helpers here can be
tested without the multi-gigabyte environment.
"""

from __future__ import annotations

import argparse
import json
import os
import sys
from collections.abc import Iterable
from pathlib import Path
from typing import Any

MODEL = "pyannote/speaker-diarization-community-1"


def turns_from_tracks(tracks: Iterable[tuple[Any, Any, str]]) -> list[dict[str, Any]]:
    """Convert pyannote (segment, track, label) triples into sorted turn dicts.

    Times are rounded to milliseconds. Segments with no length are dropped.
    """
    turns = [
        {"start": round(float(seg.start), 3), "end": round(float(seg.end), 3), "speaker": str(label)}
        for seg, _track, label in tracks
        if float(seg.end) > float(seg.start)
    ]
    return sorted(turns, key=lambda t: (t["start"], t["end"]))


def _positive_int(value: str) -> int:
    number = int(value)
    if number < 1:
        raise argparse.ArgumentTypeError("must be 1 or more")
    return number


def parse_args(argv: list[str]) -> argparse.Namespace:
    """Parse the command line. Exits with usage on bad or contradictory counts."""
    parser = argparse.ArgumentParser(description="Speaker diarization with pyannote.")
    parser.add_argument("audio")
    parser.add_argument("out")
    parser.add_argument("--speakers", type=_positive_int, help="exact number of speakers")
    parser.add_argument("--min-speakers", type=_positive_int)
    parser.add_argument("--max-speakers", type=_positive_int)
    args = parser.parse_args(argv)
    if args.speakers is not None and (args.min_speakers or args.max_speakers):
        parser.error("--speakers cannot be combined with --min-speakers/--max-speakers")
    if args.min_speakers and args.max_speakers and args.min_speakers > args.max_speakers:
        parser.error("--min-speakers cannot exceed --max-speakers")
    return args


def pipeline_kwargs(args: argparse.Namespace) -> dict[str, int]:
    """Only the speaker-count options that were actually given."""
    options = {
        "num_speakers": args.speakers,
        "min_speakers": args.min_speakers,
        "max_speakers": args.max_speakers,
    }
    return {name: value for name, value in options.items() if value is not None}


def run(args: argparse.Namespace) -> list[dict[str, Any]]:
    """Load the pipeline, diarize the audio, and return the turns."""
    # Heavy, and only installed in the service venv, so imported here on purpose.
    from pyannote.audio import Pipeline  # pyright: ignore[reportMissingImports]

    pipeline = Pipeline.from_pretrained(MODEL, token=os.environ.get("HF_TOKEN") or None)
    if pipeline is None:
        raise RuntimeError(f"could not load {MODEL}: check HF_TOKEN and that its terms are accepted")
    output = pipeline(args.audio, **pipeline_kwargs(args))
    # The exclusive variant never overlaps two speakers, which is what a
    # word-by-word merge wants. Older pipelines return the annotation itself.
    annotation = getattr(output, "exclusive_speaker_diarization", None)
    if annotation is None:
        annotation = getattr(output, "speaker_diarization", output)
    return turns_from_tracks(annotation.itertracks(yield_label=True))


def main(argv: list[str]) -> int:
    """CLI entry point. Writes OUT_JSON atomically; non-zero on any failure."""
    args = parse_args(argv)
    if not Path(args.audio).is_file():
        print(f"Error: audio file not found: {args.audio}", file=sys.stderr)
        return 1
    try:
        turns = run(args)
    except Exception as err:  # noqa: BLE001 - report any model failure and exit non-zero
        print(f"Error: diarization failed: {err}", file=sys.stderr)
        return 1
    if not turns:
        print("Error: diarization found no speech", file=sys.stderr)
        return 1
    out = Path(args.out)
    partial = out.with_name(out.name + ".partial")
    partial.write_text(json.dumps(turns), encoding="utf-8")
    partial.replace(out)
    speakers = len({t["speaker"] for t in turns})
    print(f"{len(turns)} turns, {speakers} speakers -> {out}", file=sys.stderr)
    return 0


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