#!/usr/bin/env python3
"""Transcribe an audio/video file with faster-whisper and write TXT + SRT outputs."""
from __future__ import annotations

import argparse
import json
import time
from pathlib import Path

from faster_whisper import WhisperModel


def fmt_srt(t: float) -> str:
    h = int(t // 3600)
    m = int((t % 3600) // 60)
    s = int(t % 60)
    ms = int((t - int(t)) * 1000)
    return f"{h:02d}:{m:02d}:{s:02d},{ms:03d}"


def fmt_plain(t: float) -> str:
    h = int(t // 3600)
    m = int((t % 3600) // 60)
    s = int(t % 60)
    return f"{h:02d}:{m:02d}:{s:02d}"


def main() -> None:
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument("media", help="Audio/video file to transcribe")
    ap.add_argument("--outdir", default=None, help="Output directory (default: beside media file)")
    ap.add_argument("--model", default="small", help="Whisper model size/name (default: small)")
    ap.add_argument("--language", default=None, help="Language code, e.g. ru or en. Omit for auto-detect")
    ap.add_argument("--device", default="cpu", help="Device for faster-whisper (default: cpu)")
    ap.add_argument("--compute-type", default="int8", help="Compute type (default: int8)")
    ap.add_argument("--beam-size", type=int, default=5)
    ap.add_argument("--no-vad", action="store_true", help="Disable VAD filtering")
    args = ap.parse_args()

    media = Path(args.media).expanduser().resolve()
    outdir = Path(args.outdir).expanduser().resolve() if args.outdir else media.parent
    outdir.mkdir(parents=True, exist_ok=True)
    stem = media.stem
    out_txt = outdir / f"{stem}_transcript.txt"
    out_srt = outdir / f"{stem}_transcript.srt"
    out_json = outdir / f"{stem}_segments.json"

    started = time.time()
    model = WhisperModel(args.model, device=args.device, compute_type=args.compute_type)
    segments, info = model.transcribe(
        str(media),
        language=args.language,
        vad_filter=not args.no_vad,
        beam_size=args.beam_size,
        condition_on_previous_text=True,
        word_timestamps=False,
    )

    plain_parts: list[str] = []
    srt_parts: list[str] = []
    json_parts: list[dict] = []

    for i, seg in enumerate(segments, 1):
        text = seg.text.strip()
        if not text:
            continue
        plain_parts.append(f"[{fmt_plain(seg.start)}–{fmt_plain(seg.end)}] {text}")
        srt_parts.append(f"{i}\n{fmt_srt(seg.start)} --> {fmt_srt(seg.end)}\n{text}\n")
        json_parts.append({"i": i, "start": seg.start, "end": seg.end, "text": text})
        if i % 20 == 0:
            out_txt.write_text("\n".join(plain_parts), encoding="utf-8")
            out_srt.write_text("\n".join(srt_parts), encoding="utf-8")
            out_json.write_text(json.dumps(json_parts, ensure_ascii=False, indent=2), encoding="utf-8")
            print(f"segments={i} elapsed={time.time() - started:.1f}s", flush=True)

    header = (
        f"Транскрибация аудио: {media.name}\n"
        f"Язык: {getattr(info, 'language', args.language) or 'auto'}\n"
        f"Модель: faster-whisper {args.model} / {args.device} {args.compute_type}\n\n"
    )
    out_txt.write_text(header + "\n".join(plain_parts), encoding="utf-8")
    out_srt.write_text("\n".join(srt_parts), encoding="utf-8")
    out_json.write_text(
        json.dumps({"language": getattr(info, "language", None), "duration": getattr(info, "duration", None), "segments": json_parts}, ensure_ascii=False, indent=2),
        encoding="utf-8",
    )
    print(f"DONE segments={len(json_parts)} elapsed={time.time() - started:.1f}s txt={out_txt} srt={out_srt}")


if __name__ == "__main__":
    main()