#!/usr/bin/env python3
"""Whisper-based forced alignment for TTS-generated audio.

Uses faster-whisper (preferred) or openai-whisper to produce word-level
timestamps. Outputs JSON to stdout.

Usage:
    python3 align.py <audio_path> [--text TEXT] [--language LANG] [--model MODEL]

Output format:
    {"ok": true, "words": [{"word": "...", "start": 0.0, "end": 0.5}, ...]}
    {"ok": false, "error": "error message"}
"""

import sys
import json
import argparse


def transcribe_with_faster_whisper(audio_path, language=None, model_size="base"):
    """Transcribe using faster-whisper and return word-level timestamps."""
    from faster_whisper import WhisperModel

    model = WhisperModel(model_size, device="cpu", compute_type="int8")

    segments, _info = model.transcribe(
        audio_path,
        word_timestamps=True,
        language=language,
        vad_filter=True,
    )

    words = []
    for segment in segments:
        for word_info in segment.words:
            words.append({
                "word": word_info.word.strip(),
                "start": round(word_info.start, 3),
                "end": round(word_info.end, 3),
            })

    return words


def transcribe_with_whisper(audio_path, language=None, model_size="base"):
    """Fallback: transcribe using openai-whisper and return word-level timestamps."""
    import whisper

    model = whisper.load_model(model_size)
    result = model.transcribe(
        audio_path,
        word_timestamps=True,
        language=language,
    )

    words = []
    for segment in result.get("segments", []):
        for word_info in segment.get("words", []):
            words.append({
                "word": word_info["word"].strip(),
                "start": round(word_info["start"], 3),
                "end": round(word_info["end"], 3),
            })

    return words


def main():
    parser = argparse.ArgumentParser(description="Whisper forced alignment")
    parser.add_argument("audio_path", help="Path to audio file")
    parser.add_argument("--text", default=None, help="Reference text (unused, reserved)")
    parser.add_argument("--language", default=None, help="Language hint (e.g. 'zh', 'en')")
    parser.add_argument("--model", default="base", help="Whisper model size (tiny/base/small/medium/large)")
    args = parser.parse_args()

    try:
        try:
            words = transcribe_with_faster_whisper(
                args.audio_path,
                language=args.language,
                model_size=args.model,
            )
        except ImportError:
            words = transcribe_with_whisper(
                args.audio_path,
                language=args.language,
                model_size=args.model,
            )

        print(json.dumps({"ok": True, "words": words}))
    except Exception as e:
        print(json.dumps({"ok": False, "error": str(e)}))
        sys.exit(0)


if __name__ == "__main__":
    main()