import gradio as gr
import whisperx
import torch
import json
import numpy as np

device = "cuda" if torch.cuda.is_available() else "cpu"

def align_words(audio_file, transcript):
    audio = whisperx.load_audio(audio_file)

    # ── Trim leading silence ──────────────────────────────────────────────────
    FRAME  = 160      # 10ms at 16kHz
    THRESH = 0.001    # RMS energy threshold
    speech_start = 0
    for i in range(0, len(audio) - FRAME, FRAME):
        if float(np.sqrt(np.mean(audio[i:i+FRAME]**2))) > THRESH:
            speech_start = i
            break
    audio_trimmed = np.ascontiguousarray(audio[speech_start:], dtype=np.float32)

    # ── Forced alignment ──────────────────────────────────────────────────────
    segments = [{"text": transcript, "start": 0, "end": 9999}]
    model_a, metadata = whisperx.load_align_model(language_code="en", device=device)
    result = whisperx.align(segments, model_a, metadata, audio_trimmed, device)

    # ── Build output ──────────────────────────────────────────────────────────
    output = []
    for w in result["word_segments"]:
        start = w["start"]
        end   = w["end"]

        # WhisperX sometimes returns ms instead of seconds — normalise
        if start > 500:
            start = round(start / 1000, 3)
            end   = round(end   / 1000, 3)
        else:
            start = round(start, 3)
            end   = round(end,   3)

        output.append({"word": w["word"], "start": start, "end": end})

    return json.dumps(output, indent=2)

demo = gr.Interface(
    fn=align_words,
    inputs=[
        gr.Audio(type="filepath", label="Audio"),
        gr.Textbox(lines=10, label="Exact Transcript")
    ],
    outputs=gr.Textbox(label="Word Timestamps JSON"),
    title="Fast Word Timestamp Alignment"
)
demo.launch()