#!/usr/bin/env python3
"""Local transcription with faster-whisper.

Tuned for NVIDIA GTX 1660 Ti (Turing GTX 16xx): uses int8_float32 because
float16 / int8_float16 are broken/slow on this GPU family.

Outputs (next to input, same basename):
  <name>.txt   plain text, one line per segment
  <name>.srt   subtitles with timestamps
  <name>.json  raw segments (start, end, text) for later processing

Raw ASR only. Italian-dialect adaptation/cleanup is a separate step.
"""
import argparse
import json
import sys
from pathlib import Path

from faster_whisper import WhisperModel, BatchedInferencePipeline

# Bias recognition toward the domain vocabulary (corporate networks).
DEFAULT_PROMPT = (
    "Lezione tecnica su reti aziendali. Termini ricorrenti: rete, server, "
    "switch, router, firewall, VLAN, indirizzo IP, DNS, DHCP, VPN, gateway, "
    "backbone, cablaggio, dominio, Active Directory, backup, NAS, subnet, "
    "porta, patch, armadio rack, centralino, PLC."
)


def fmt_ts(seconds: float) -> str:
    """seconds -> SRT timestamp HH:MM:SS,mmm"""
    ms = int(round(seconds * 1000))
    h, ms = divmod(ms, 3_600_000)
    m, ms = divmod(ms, 60_000)
    s, ms = divmod(ms, 1000)
    return f"{h:02d}:{m:02d}:{s:02d},{ms:03d}"


def main() -> int:
    ap = argparse.ArgumentParser(description="Local faster-whisper transcription")
    ap.add_argument("audio", type=Path, help="input audio/video file")
    ap.add_argument("--model", default="large-v3-turbo",
                    help="model name or path (default: large-v3-turbo)")
    ap.add_argument("--device", default="cuda", choices=["cuda", "cpu"])
    ap.add_argument("--compute-type", default="int8_float32",
                    help="int8_float32 for GTX 1660 Ti (NOT float16)")
    ap.add_argument("--language", default="it")
    ap.add_argument("--task", default="transcribe", choices=["transcribe", "translate"])
    ap.add_argument("--beam-size", type=int, default=5)
    ap.add_argument("--prompt", default=DEFAULT_PROMPT,
                    help="initial_prompt to bias vocabulary")
    ap.add_argument("--no-vad", action="store_true", help="disable Silero VAD filter")
    ap.add_argument("--batched", action="store_true",
                    help="use BatchedInferencePipeline: independent chunks, "
                         "no cross-context (kills hallucination loops)")
    ap.add_argument("--batch-size", type=int, default=8,
                    help="batch size for --batched (default 8)")
    ap.add_argument("--outdir", type=Path, default=None,
                    help="output dir (default: alongside input)")
    args = ap.parse_args()

    if not args.audio.is_file():
        print(f"ERROR: file not found: {args.audio}", file=sys.stderr)
        return 1

    outdir = args.outdir or args.audio.parent
    outdir.mkdir(parents=True, exist_ok=True)
    stem = args.audio.stem
    txt_path = outdir / f"{stem}.txt"
    srt_path = outdir / f"{stem}.srt"
    json_path = outdir / f"{stem}.json"

    print(f"[load] model={args.model} device={args.device} "
          f"compute_type={args.compute_type}", flush=True)
    try:
        model = WhisperModel(args.model, device=args.device,
                             compute_type=args.compute_type)
    except Exception as e:  # noqa: BLE001
        print(f"[warn] GPU load failed ({e}); falling back to CPU int8", flush=True)
        model = WhisperModel(args.model, device="cpu", compute_type="int8")

    print(f"[run] transcribing: {args.audio.name} "
          f"(batched={args.batched})", flush=True)
    if args.batched:
        # Independent chunks: no condition_on_previous_text, VAD forced on.
        # Best for degraded/dialect audio prone to repetition loops.
        batched = BatchedInferencePipeline(model=model)
        segments, info = batched.transcribe(
            str(args.audio),
            language=args.language,
            task=args.task,
            beam_size=args.beam_size,
            initial_prompt=args.prompt,
            batch_size=args.batch_size,
            repetition_penalty=1.2,
        )
    else:
        segments, info = model.transcribe(
            str(args.audio),
            language=args.language,
            task=args.task,
            beam_size=args.beam_size,
            initial_prompt=args.prompt,
            vad_filter=not args.no_vad,
            vad_parameters=dict(min_silence_duration_ms=500),
            condition_on_previous_text=True,
        )
    print(f"[info] detected language={info.language} "
          f"prob={info.language_probability:.2f} "
          f"duration={info.duration:.1f}s", flush=True)

    seg_records = []
    with txt_path.open("w", encoding="utf-8") as ftxt, \
         srt_path.open("w", encoding="utf-8") as fsrt:
        for i, seg in enumerate(segments, 1):
            text = seg.text.strip()
            # stream progress to stdout
            print(f"  [{fmt_ts(seg.start)} -> {fmt_ts(seg.end)}] {text}", flush=True)
            ftxt.write(text + "\n")
            fsrt.write(f"{i}\n{fmt_ts(seg.start)} --> {fmt_ts(seg.end)}\n{text}\n\n")
            seg_records.append({"id": i, "start": seg.start,
                                "end": seg.end, "text": text})

    json_path.write_text(
        json.dumps({"language": info.language,
                    "duration": info.duration,
                    "model": args.model,
                    "segments": seg_records},
                   ensure_ascii=False, indent=2),
        encoding="utf-8")

    print(f"\n[done] {len(seg_records)} segments", flush=True)
    print(f"  txt : {txt_path}")
    print(f"  srt : {srt_path}")
    print(f"  json: {json_path}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())