import gradio as gr
import json

from faster_whisper import WhisperModel

# Fast model
model = WhisperModel(
    "base",
    device="cpu",
    compute_type="int8"
)


def sec_to_srt(sec):
    total_ms = int(round(sec * 1000))

    hours = total_ms // 3600000
    total_ms %= 3600000

    minutes = total_ms // 60000
    total_ms %= 60000

    seconds = total_ms // 1000
    milliseconds = total_ms % 1000

    return (
        f"{hours:02d}:"
        f"{minutes:02d}:"
        f"{seconds:02d},"
        f"{milliseconds:03d}"
    )


def generate_srt(audio_file, timestamp_json, language):

    if audio_file is None:
        return "Please upload audio"

    try:
        timestamps = json.loads(timestamp_json)
    except Exception as e:
        return f"Invalid JSON:\n{e}"

    # Transcribe ONCE
    segments, info = model.transcribe(
        audio_file,
        language=language.strip(),
        word_timestamps=True,
        beam_size=1,
        vad_filter=False
    )

    # Collect all words
    all_words = []

    for segment in segments:

        if segment.words is None:
            continue

        for word in segment.words:

            if word.start is None or word.end is None:
                continue

            all_words.append({
                "text": word.word.strip(),
                "start": float(word.start),
                "end": float(word.end)
            })

    # Generate SRT
    srt_blocks = []

    for idx, item in enumerate(timestamps, start=1):

        start = float(item["start"])
        end = float(item["end"])

        words = []

        for w in all_words:

            center = (w["start"] + w["end"]) / 2

            if start <= center <= end:
                words.append(w["text"])

        text = " ".join(words).strip()

        srt_blocks.append(
            f"{idx}\n"
            f"{sec_to_srt(start)} --> {sec_to_srt(end)}\n"
            f"{text}\n"
        )

    return "\n".join(srt_blocks)


demo = gr.Interface(
    fn=generate_srt,
    inputs=[
        gr.Audio(
            type="filepath",
            label="Audio File"
        ),
        gr.Textbox(
            lines=20,
            label="Timestamp JSON"
        ),
        gr.Textbox(
            value="en",
            label="Language Code"
        )
    ],
    outputs=gr.Textbox(
        lines=30,
        label="Generated SRT"
    ),
    title="Timestamp JSON → SRT (FastWhisper Turbo)",
    description="Upload audio, paste timestamp JSON, select language, get SRT."
)

demo.launch()