# SPDX-License-Identifier: MIT
"""Transcribe a local audio file with NepTrans v4 and faster-whisper."""
from __future__ import annotations

import argparse
import json
from pathlib import Path

from faster_whisper import WhisperModel


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("audio", type=Path)
    parser.add_argument("--model", default="DontHash/neptrans-v4-ct2")
    parser.add_argument("--revision", default="v4.0.0")
    parser.add_argument("--device", choices=["cpu", "cuda"], default="cpu")
    parser.add_argument("--compute-type", default=None)
    parser.add_argument("--language", choices=["ne", "en"], default=None)
    parser.add_argument("--beam-size", type=int, default=1)
    parser.add_argument("--local-files-only", action="store_true")
    parser.add_argument("--output", type=Path)
    parser.add_argument("--json-output", type=Path)
    args = parser.parse_args()
    if not args.audio.is_file():
        parser.error(f"Audio file not found: {args.audio}")
    if args.beam_size < 1:
        parser.error("--beam-size must be positive")
    if args.output and args.output.resolve() == args.audio.resolve():
        parser.error("Output must not overwrite the input audio")
    if args.json_output and args.json_output.resolve() == args.audio.resolve():
        parser.error("JSON output must not overwrite the input audio")
    if args.output and args.json_output and args.output.resolve() == args.json_output.resolve():
        parser.error("Text and JSON outputs must be different files")
    compute_type = args.compute_type or ("int8_float16" if args.device == "cuda" else "int8")
    model = WhisperModel(
        args.model, revision=args.revision, device=args.device,
        compute_type=compute_type, local_files_only=args.local_files_only,
    )
    segments, info = model.transcribe(
        str(args.audio), language=args.language, task="transcribe",
        beam_size=args.beam_size, condition_on_previous_text=False,
    )
    rows = [{"start": s.start, "end": s.end, "text": s.text.strip()} for s in segments]
    text = "\n".join(row["text"] for row in rows) + "\n"
    if args.output:
        args.output.write_text(text, encoding="utf-8")
    else:
        print(text, end="")
    if args.json_output:
        args.json_output.write_text(json.dumps({
            "language": info.language,
            "language_probability": info.language_probability,
            "duration": info.duration,
            "segments": rows,
        }, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")


if __name__ == "__main__":
    main()