#!/usr/bin/env python3
"""Create resumable word-timestamp ASR sidecars for audiobook tracks."""

from __future__ import annotations

import argparse
import hashlib
import json
import sys
import time
from pathlib import Path


def sha256(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as source:
        for block in iter(lambda: source.read(1024 * 1024), b""):
            digest.update(block)
    return digest.hexdigest()


def serialize_segment(segment) -> dict[str, object]:
    return {
        "id": segment.id,
        "start": round(segment.start, 3),
        "end": round(segment.end, 3),
        "text": segment.text.strip(),
        "avg_logprob": round(segment.avg_logprob, 5),
        "compression_ratio": round(segment.compression_ratio, 5),
        "no_speech_prob": round(segment.no_speech_prob, 5),
        "words": [
            {
                "start": round(word.start, 3),
                "end": round(word.end, 3),
                "word": word.word,
                "probability": round(word.probability, 5),
            }
            for word in (segment.words or ())
        ],
    }


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("--source-dir", type=Path, required=True)
    parser.add_argument("--output-dir", type=Path, required=True)
    parser.add_argument("--model", default="base")
    parser.add_argument("--model-cache", type=Path, required=True)
    parser.add_argument("--max-files", type=int)
    parser.add_argument("--force", action="store_true")
    args = parser.parse_args()

    try:
        from faster_whisper import WhisperModel
    except ImportError as error:
        raise SystemExit("faster-whisper is missing from PYTHONPATH") from error

    files = sorted(args.source_dir.glob("*.mp3"))
    if args.max_files is not None:
        files = files[: max(0, args.max_files)]
    if not files:
        parser.error("No MP3 files found")
    args.output_dir.mkdir(parents=True, exist_ok=True)
    model = WhisperModel(
        args.model,
        device="cpu",
        compute_type="int8",
        download_root=str(args.model_cache),
        local_files_only=True,
    )

    completed = 0
    for index, audio_path in enumerate(files, 1):
        output_path = args.output_dir / f"{audio_path.stem[:4]}.json"
        if output_path.is_file() and not args.force:
            print(f"[{index}/{len(files)}] skip {audio_path.name}", flush=True)
            completed += 1
            continue
        started = time.monotonic()
        print(f"[{index}/{len(files)}] transcribe {audio_path.name}", flush=True)
        segments_iter, info = model.transcribe(
            str(audio_path),
            language="ru",
            beam_size=5,
            vad_filter=True,
            word_timestamps=True,
            condition_on_previous_text=True,
        )
        segments = [serialize_segment(segment) for segment in segments_iter]
        payload = {
            "schema": 1,
            "source": str(audio_path.resolve()),
            "source_bytes": audio_path.stat().st_size,
            "source_sha256": sha256(audio_path),
            "language": info.language,
            "language_probability": round(info.language_probability, 5),
            "duration": round(info.duration, 3),
            "duration_after_vad": round(info.duration_after_vad, 3),
            "elapsed_seconds": round(time.monotonic() - started, 3),
            "segments": segments,
        }
        temporary = output_path.with_suffix(".json.tmp")
        temporary.write_text(json.dumps(payload, ensure_ascii=False, indent=2), encoding="utf-8")
        temporary.replace(output_path)
        completed += 1
        print(
            f"[{index}/{len(files)}] wrote {output_path.name}: "
            f"{len(segments)} segments in {payload['elapsed_seconds']}s",
            flush=True,
        )
    print(json.dumps({"files": len(files), "completed": completed}, ensure_ascii=False))
    return 0


if __name__ == "__main__":
    sys.exit(main())