# Copyright (c) 2026 The Anuvani-Diarization authors
# SPDX-License-Identifier: MIT OR Apache-2.0
"""Diarization in one line, at full speed.

    from anuvani_diar import Diarizer
    diarize = Diarizer()                          # every visible GPU (or the CPU); Diarizer(devices=2): two GPUs
    result = diarize("meeting.wav")               # -> Diarization
    result.save_rttm("meeting.rttm")
    for r in diarize.run(["a.wav", "b.flac"]):    # many files; results arrive as each one finishes
        print(r.uri, r.speakers)

Everything except the two networks runs in the Rust engine (`anuvani_diar._core`), on its own threads: decoding,
windowing, batching across all GPUs, masks, clustering, the decoder and the RTTM text. The output is the release's
(the original program, `reference/diarize.py` in the model repository): byte-identical on the setup the release's
outputs were produced with, and near-identical, checked at start-up, on any other GPU or PyTorch version.
"""
import json
import os
from dataclasses import dataclass, field
from pathlib import Path
from typing import Iterable, Iterator, List, Optional, Tuple, Union

from . import _core

_HERE = Path(__file__).resolve().parent


def default_bundle() -> Path:
    """`$ANUVANI_DIAR_BUNDLE`; else, in a source tree, the bundle exported beside it (artifacts/bundle, made by
    tools/export_bundle.py from the release's checkpoints: safetensors, the Indic model only); else the model folder of
    the Hugging Face repository, downloaded once and sha256-checked (anuvani_diar.hub)."""
    env = os.environ.get("ANUVANI_DIAR_BUNDLE")
    if env:
        return Path(env)
    p = _HERE.parents[1] / "artifacts" / "bundle"
    if (p / "MANIFEST.json").is_file():
        return p
    from .hub import bundle_dir
    return bundle_dir()


def devices_arg(text: str) -> Union[str, int]:
    """A --devices value as `Diarizer(devices=...)` takes it: a bare number is a count of GPUs."""
    text = text.strip()
    return int(text) if text.isdigit() else text


def default_golden() -> Optional[Path]:
    env = os.environ.get("ANUVANI_DIAR_GOLDEN")
    if env:
        return Path(env)
    from .hub import assets_dir
    p = assets_dir() / "nn_golden.json"
    return p if p.is_file() else None


@dataclass
class Diarization:
    """One file's result. `rttm` is the exact text the release writes; `segments` are read back from it."""
    uri: str
    path: str
    rttm: Optional[str]
    error: Optional[str] = None
    duration_s: Optional[float] = None
    timing: dict = field(default_factory=dict)
    record: Optional[dict] = None
    index: int = -1

    @property
    def ok(self) -> bool:
        return self.error is None and self.rttm is not None

    @property
    def segments(self) -> List[Tuple[float, float, str]]:
        """(start, end, speaker) per RTTM line, in the file's order."""
        out = []
        for line in (self.rttm or "").splitlines():
            f = line.split()
            out.append((float(f[3]), float(f[3]) + float(f[4]), f[7]))
        return out

    @property
    def speakers(self) -> List[str]:
        return sorted({s for _, _, s in self.segments})

    def save_rttm(self, path: Union[str, Path]) -> Path:
        if not self.ok:
            raise RuntimeError(self.error or f"no output for {self.path}")
        path = Path(path)
        path.write_bytes(self.rttm.encode())
        return path


class Diarizer:
    """The release's diarization.

    devices: where the networks run. "auto" (default): every visible GPU, or the CPU when there is none; a number N:
    the first N GPUs; a list of GPU indices or device names ([0, 3], ["cuda:1"], "0,3"); or "cpu". The engine
    spreads every file's windows over all of them.
    workers_per_device: kept at 1; one CUDA stream per GPU is what keeps the outputs deterministic.
    precision: "exact" (default: the release's FP32 networks and settings), or a faster reduced-precision preset whose
    outputs differ slightly: "ffn16", "tf32conv" and "mixed"; "fp32", "tf32" and "fp16" are for measurement only.
    strict: byte-identical outputs need the setup the release's outputs were produced with. On any other GPU or
    PyTorch version, with strict=False (default), the exact lane runs if its start-up self-test stays within
    near-identical bounds of the release's outputs (a warning says so, and `info()["exactness"]` gives the numbers);
    strict=True refuses to start instead.
    """

    def __init__(self, bundle: Union[str, Path, None] = None, devices: Union[str, int, List, None] = "auto",
                 workers_per_device: int = 1, threads: int = 2, repack: bool = True, keep_records: bool = False,
                 decode_threads: int = 4, post_threads: Optional[int] = None, max_ready: Optional[int] = None, golden=None,
                 precision: str = "exact", strict: bool = False):
        from .nn.torch_backend import TorchNN
        self.bundle = Path(bundle) if bundle else default_bundle()
        manifest = json.loads((self.bundle / "MANIFEST.json").read_text())
        self.manifest = manifest
        model = manifest["modes"]["indic"]
        if golden is None:
            golden = default_golden()
            if golden is None:
                raise FileNotFoundError(
                    "assets/nn_golden.json not found: every GPU checks the networks against the release's own outputs "
                    "before use. Set ANUVANI_DIAR_GOLDEN to the file (or pass golden=False to run without the check).")
        elif golden is False:
            golden = None
        self.nn = TorchNN(self.bundle, devices, workers_per_device, threads, golden=golden, precision=precision,
                          strict=strict)
        self.engine = _core.Engine(
            self.nn, clustering=model["clustering"], decoder_path=str(self.bundle / model["decoder"]),
            decode_threads=decode_threads, post_threads=post_threads or max(4, (os.cpu_count() or 8) // 2),
            repack=repack, max_ready=max_ready, keep_records=keep_records)

    def __call__(self, path: Union[str, Path]) -> Diarization:
        return next(iter(self.run([path])))

    def run(self, paths: Iterable[Union[str, Path]], out_dir: Union[str, Path, None] = None) -> Iterator[Diarization]:
        """Diarise many files; yields results in completion order. With `out_dir`, also writes `<stem>.rttm` there."""
        paths = [Path(p) for p in paths]
        jobs = [(p.stem, str(p), str(p)) for p in paths]
        out = Path(out_dir) if out_dir else None
        for r in self.engine.run(jobs):
            d = Diarization(uri=r["uri"], path=r["name"], rttm=r["rttm"], error=r["error"], duration_s=r["duration_s"],
                            timing=r["timing"], record=r.get("record"), index=r["index"])
            if out is not None and d.ok:
                out.mkdir(parents=True, exist_ok=True)
                d.save_rttm(out / f"{d.uri}.rttm")
            yield d

    def info(self) -> dict:
        return dict(self.nn.info(), version=self.manifest["version"], bundle=str(self.bundle))