# Copyright (c) 2026 The Anuvani-Diarization authors
# SPDX-License-Identifier: MIT OR Apache-2.0
"""anuvani-diar: the release's `diarize.py --mode indic` command line, on the Rust engine.

  anuvani-diar --out OUT_DIR audio1.wav [audio2.flac ...]
  anuvani-diar --out OUT_DIR --in-dir AUDIO_DIR

RTTM files byte-identical to release/anuvani_diar_v1/diarize.py --mode indic, and its other flags. Differences: there
is one model, so no --mode; run.json names the model's settings `model_config`; every visible GPU is used unless
--devices says otherwise; files finish in any order (run.json keeps the input order); a file that fails gets an "error"
entry instead of stopping the run, and the exit status is 1.
"""
import argparse
import hashlib
import json
import platform
import sys
from pathlib import Path

AUDIO_EXT = ("*.wav", "*.flac", "*.mp3", "*.m4a", "*.ogg")


def main(argv=None):
    p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    p.add_argument("audio", nargs="*", type=Path)
    p.add_argument("--in-dir", type=Path)
    p.add_argument("--out", type=Path, required=True)
    p.add_argument("--bundle", type=Path, default=None)
    p.add_argument("--devices", "--device", dest="devices", default="auto", help="which devices run the networks: auto (every visible GPU, else the CPU), a number of GPUs (2: the first two), GPU indices (0,3), cuda:N, or cpu")
    p.add_argument("--batch-size", type=int, default=8)
    p.add_argument("--threads", type=int, default=2)
    a = p.parse_args(argv)
    if a.batch_size != 8:
        raise SystemExit("--batch-size: only 8, the release's verified setting, is supported")
    files = list(a.audio)
    if a.in_dir:
        for ext in AUDIO_EXT:
            files.extend(sorted(a.in_dir.glob(ext)))
    files = [f for f in files if f.is_file()]
    if not files:
        raise SystemExit("no audio files given")
    from .api import Diarizer, devices_arg
    d = Diarizer(bundle=a.bundle, devices=devices_arg(a.devices), threads=a.threads)
    manifest = d.manifest
    a.out.mkdir(parents=True, exist_ok=True)
    log = dict(version=manifest["version"], model_config=manifest["modes"]["indic"], bundle_sha256=manifest["files"],
               devices=a.devices, batch_size=a.batch_size, threads=a.threads, python=platform.python_version(), files=[])
    recs, failed = [None] * len(files), 0
    for r in d.run(files):
        f = files[r.index]
        if r.ok:
            out = r.save_rttm(a.out / f"{r.uri}.rttm")
            rec = dict(input=str(f), rttm=str(out), duration_s=round(r.duration_s, 3), seconds=round(r.timing["wall_s"], 2),
                       rttm_sha256=hashlib.sha256(out.read_bytes()).hexdigest())
        else:
            failed += 1
            rec = dict(input=str(f), error=r.error)
        recs[r.index] = rec
        print(json.dumps(rec), flush=True)
    log["files"] = recs
    (a.out / "run.json").write_text(json.dumps(log, indent=1))
    return 1 if failed else 0


if __name__ == "__main__":
    sys.exit(main())