#!/usr/bin/env python3
"""Transcribe a Khanega lecture audio file via OpenAI Whisper (Persian).

Whisper API has a 25 MB file size limit per request. For lectures > 25 MB (typical
1-hour lecture is 30-60 MB as mp3), we chunk the audio into ~10-min segments using
pydub, transcribe each, and stitch the segments back together with adjusted timestamps.

Output: <input>.fa.json — list of {start, end, text} segments in Persian.

Usage:
    python transcribe.py audio/<file>.mp3
    python transcribe.py audio/<file>.mp3 --chunk-min 8   # smaller chunks if needed
"""
import argparse
import json
import math
import os
import shutil
import subprocess
import sys
import tempfile
from pathlib import Path

from openai import OpenAI


def ffprobe_duration_s(src: Path) -> float:
    """Return audio duration in seconds via ffprobe."""
    r = subprocess.run(
        ["ffprobe", "-v", "error", "-show_entries", "format=duration",
         "-of", "default=nokey=1:noprint_wrappers=1", str(src)],
        capture_output=True, text=True, check=True,
    )
    return float(r.stdout.strip())


def chunk_audio(src: Path, chunk_s: int):
    """Yield (chunk_path, offset_s) tuples. Uses ffmpeg to extract 64kbps mp3 chunks
    (small enough that 10-min chunks stay well under Whisper's 25 MB cap)."""
    total_s = ffprobe_duration_s(src)
    tmp = Path(tempfile.mkdtemp(prefix="khanega_chunks_"))
    n_chunks = math.ceil(total_s / chunk_s)
    for i in range(n_chunks):
        offset = i * chunk_s
        chunk_path = tmp / f"chunk_{i:03d}.mp3"
        subprocess.run(
            ["ffmpeg", "-y", "-loglevel", "error",
             "-ss", str(offset), "-t", str(chunk_s),
             "-i", str(src),
             "-vn", "-ac", "1", "-ar", "16000", "-b:a", "64k",
             str(chunk_path)],
            check=True,
        )
        yield chunk_path, float(offset)


def transcribe_chunk(client: OpenAI, chunk_path: Path) -> list[dict]:
    """Transcribe one chunk; returns list of {start, end, text} relative to chunk start."""
    with chunk_path.open("rb") as f:
        resp = client.audio.transcriptions.create(
            model="whisper-1",
            file=f,
            language="fa",
            response_format="verbose_json",
            timestamp_granularities=["segment"],
        )
    out = []
    for seg in resp.segments:
        out.append({
            "start": float(seg.start),
            "end": float(seg.end),
            "text": seg.text.strip(),
        })
    return out


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("audio_file", type=Path)
    ap.add_argument("--chunk-min", type=int, default=10, help="chunk length in minutes (default 10)")
    ap.add_argument("--force", action="store_true", help="re-transcribe even if output exists")
    args = ap.parse_args()

    audio_path = args.audio_file
    if not audio_path.exists():
        sys.exit(f"file not found: {audio_path}")

    out_path = audio_path.with_suffix(audio_path.suffix + ".fa.json")
    if out_path.exists() and not args.force:
        print(f"already transcribed → {out_path} (use --force to redo)")
        return

    client = OpenAI(api_key=os.environ["OPENAI_API_KEY"])
    chunk_s = args.chunk_min * 60

    all_segments = []
    tmp_dirs = []  # so we can clean at the end
    for chunk_path, offset_s in chunk_audio(audio_path, chunk_s):
        tmp_dirs.append(chunk_path.parent)
        print(f"  transcribing {chunk_path.name} (offset {offset_s:.1f}s)...", flush=True)
        segments = transcribe_chunk(client, chunk_path)
        for seg in segments:
            all_segments.append({
                "start": round(seg["start"] + offset_s, 2),
                "end": round(seg["end"] + offset_s, 2),
                "text": seg["text"],
            })

    # Cleanup chunk dirs (set in case multiple iters returned same dir)
    for d in set(tmp_dirs):
        shutil.rmtree(d, ignore_errors=True)

    out_path.write_text(json.dumps({
        "source": str(audio_path),
        "language": "fa",
        "model": "whisper-1",
        "segments": all_segments,
    }, ensure_ascii=False, indent=2))
    out_path.chmod(0o600)
    print(f"\n✓ {len(all_segments)} segments → {out_path}")


if __name__ == "__main__":
    main()