#!/opt/homebrew/bin/python3.11
import sys
import json
import os
import argparse
from faster_whisper import WhisperModel

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from redis_publisher import RedisPublisher


def run_asr(video_path, output_path, uuid: str = ""):
    publisher = RedisPublisher(uuid) if uuid else None
    if publisher:
        publisher.info("asr", "ASR_START")

    if publisher:
        publisher.info("asr", "Loading Whisper model...")

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

    if publisher:
        publisher.info("asr", f"Transcribing: {video_path}")

    segments, info = model.transcribe(video_path, beam_size=5)

    if publisher:
        publisher.info("asr", f"ASR_LANGUAGE:{info.language}")

    results = []
    total_segments = 0

    for segment in segments:
        results.append(
            {"start": segment.start, "end": segment.end, "text": segment.text.strip()}
        )
        total_segments += 1
        if total_segments % 100 == 0:
            if publisher:
                publisher.progress(
                    "asr", total_segments, 0, f"Segment {total_segments}"
                )

    output = {
        "language": info.language,
        "language_probability": info.language_probability,
        "segments": results,
    }

    with open(output_path, "w") as f:
        json.dump(output, f, indent=2)

    if publisher:
        publisher.complete("asr", f"{len(results)} segments")


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="ASR Transcription")
    parser.add_argument("video_path", help="Path to video file")
    parser.add_argument("output_path", help="Output JSON path")
    parser.add_argument("--uuid", "-u", help="UUID for Redis progress", default="")
    args = parser.parse_args()

    run_asr(args.video_path, args.output_path, args.uuid)