"""Whisper transcription with segment timestamps.

Reads audio/*.mp3, writes transcripts/<slug>.json:
{
  "guid": ..., "ep_num": ..., "title": ..., "pub_date": ...,
  "audio_url": ..., "model": ..., "segments": [
     {"start": 12.34, "end": 18.90, "text": "..."}, ...
  ]
}
Fully resumable: finished transcripts are skipped unless --force.
Use --minutes N to transcribe only the first N minutes (testing).
"""
import argparse
import json
import os
import sys
import time

from config import (TRANSCRIPT_DIR, WHISPER_MODEL, DEVICE, COMPUTE_TYPE,
                    LANGUAGE, VAD_FILTER, BEAM_SIZE, WORD_TIMESTAMPS,
                    CONDITION_ON_PREVIOUS_TEXT, KEEP_AUDIO)
from feed import load_episodes
from download import audio_path, slug as ep_slug
from preprocess import prepare_audio, prepared_path, load_pcm
from lexicon import build_lexicon, tags_for_text

# faster-whisper runs one giant FFT over the whole input; multi-hour episodes
# blow up RAM (~2.5 GB arrays). Transcribe in fixed windows instead.
CHUNK_SECONDS = 600


def transcript_path(ep: dict) -> str:
    return os.path.join(TRANSCRIPT_DIR, ep_slug(ep) + ".json")


def transcribe_episode(model, ep: dict, minutes: float | None) -> dict:
    src = audio_path(ep)
    prep = prepared_path(ep)
    t0 = time.time()
    prepare_audio(src, prep, max_seconds=minutes * 60 + 5 if minutes else None)
    print(f"  prepped audio in {time.time() - t0:.0f}s -> {os.path.basename(prep)}",
          flush=True)

    pcm = load_pcm(prep)
    total_samples = len(pcm)
    if minutes:
        total_samples = min(total_samples, int(minutes * 60 * 16000))
    duration = total_samples / 16000.0

    segments = []
    n_chunks = (total_samples + CHUNK_SECONDS * 16000 - 1) // (CHUNK_SECONDS * 16000)
    for ci in range(n_chunks):
        a = ci * CHUNK_SECONDS * 16000
        b = min(a + CHUNK_SECONDS * 16000, total_samples)
        offset = a / 16000.0
        chunk = pcm[a:b]
        segs_iter, info = model.transcribe(
            chunk,
            language=LANGUAGE,
            vad_filter=VAD_FILTER,
            beam_size=BEAM_SIZE,
            word_timestamps=WORD_TIMESTAMPS,
            condition_on_previous_text=CONDITION_ON_PREVIOUS_TEXT,
        )
        for seg in segs_iter:
            text = seg.text.strip()
            if not text:
                continue
            segments.append({
                "start": round(seg.start + offset, 2),
                "end": round(seg.end + offset, 2),
                "text": text,
            })
        print(f"  chunk {ci + 1}/{n_chunks} done "
              f"(audio {(b - a) / 16000 / 60:.0f}m, "
              f"{len(segments)} segs so far)", flush=True)

    return {
        "guid": ep["guid"],
        "ep_num": ep["ep_num"],
        "title": ep["title"],
        "pub_date": ep["pub_date"],
        "audio_url": ep["audio_url"],
        "page_link": ep.get("page_link", ""),
        "duration_sec": duration,
        "model": WHISPER_MODEL,
        "language": LANGUAGE,
        "partial_minutes": minutes,
        "segments": segments,
    }


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("eps", nargs="*", type=int,
                    help="episode numbers; default: all")
    ap.add_argument("--minutes", type=float, default=None,
                    help="only transcribe first N minutes (testing)")
    ap.add_argument("--force", action="store_true",
                    help="re-transcribe even if transcript exists")
    args = ap.parse_args()

    from faster_whisper import WhisperModel

    eps = load_episodes()
    todo = []
    for ep in eps:
        if ep["ep_num"] is not None and args.eps and ep["ep_num"] not in args.eps:
            continue
        tp = transcript_path(ep)
        if os.path.exists(tp) and not args.force and not args.minutes:
            continue
        if not os.path.exists(audio_path(ep)):
            print(f"skip (no audio): {ep['title']}", file=sys.stderr)
            continue
        todo.append(ep)

    model = None  # lazy-load after we know there's work
    done = 0
    for i, ep in enumerate(todo, 1):
        label = f"Ep {ep['ep_num']}" if ep["ep_num"] else ep["title"][:40]
        if model is None:
            print(f"loading whisper model '{WHISPER_MODEL}' "
                  f"(device={DEVICE}, compute={COMPUTE_TYPE}) ...")
            model = WhisperModel(WHISPER_MODEL, device=DEVICE,
                                 compute_type=COMPUTE_TYPE)
        t0 = time.time()
        src = audio_path(ep)
        prep = prepared_path(ep)
        ok = False
        try:
            result = transcribe_episode(model, ep, args.minutes)
            outp = transcript_path(ep)
            with open(outp + ".tmp", "w", encoding="utf-8") as f:
                json.dump(result, f, ensure_ascii=False)
            os.replace(outp + ".tmp", outp)
            dur = result["duration_sec"] or 1
            rate = dur / max(time.time() - t0, 0.001)
            n_seg = len(result["segments"])
            lexicon = build_lexicon()
            ep_tag_set: set[str] = set()
            for seg in result["segments"]:
                ep_tag_set.update(tags_for_text(seg["text"], lexicon))
            print(f"[{i}/{len(todo)}] {label}: {n_seg} segments, "
                  f"{dur/60:.0f} min audio, {rate:.1f}x realtime")
            print(f"    TAGS: {' '.join('#' + t for t in sorted(ep_tag_set)) or '(none)'}")
            done += 1
            ok = True
        except Exception as e:  # noqa: BLE001
            print(f"{label}: TRANSCRIBE FAILED - {e}", file=sys.stderr)
            continue
        finally:
            # reclaim disk: prepped wav always goes; mp3 unless KEEP_AUDIO
            if os.path.exists(prep):
                try:
                    os.remove(prep)
                except OSError:
                    pass
            if not KEEP_AUDIO and ok and not args.minutes:
                try:
                    os.remove(src)
                except OSError:
                    pass

    print(f"transcribed={done} skipped_existing="
          f"{len([e for e in eps if os.path.exists(transcript_path(e))])}")


if __name__ == "__main__":
    main()