import os
import gradio as gr
import whisper
import tempfile

# Cache models so they don’t reload each time
loaded_models = {}

def load_model(model_name):
    if model_name not in loaded_models:
        loaded_models[model_name] = whisper.load_model(model_name)
    return loaded_models[model_name]

def transcribe(video, model_name, language):
    # Save uploaded video to temp file
    temp_file = tempfile.NamedTemporaryFile(delete=False, suffix=".mp4")
    temp_file.write(video)
    temp_file.close()

    # Load selected model
    model = load_model(model_name)

    # Transcribe with Whisper
    result = model.transcribe(temp_file.name, language=language, fp16=False)

    # Generate SRT text
    srt_text = []
    for i, segment in enumerate(result["segments"], start=1):
        start = segment["start"]
        end = segment["end"]
        text = segment["text"]

        srt_text.append(f"{i}")
        srt_text.append(f"{format_timestamp(start)} --> {format_timestamp(end)}")
        srt_text.append(text.strip())
        srt_text.append("")

    srt_output = "\n".join(srt_text)

    # Save as .srt file
    srt_file_path = tempfile.NamedTemporaryFile(delete=False, suffix=".srt").name
    with open(srt_file_path, "w", encoding="utf-8") as f:
        f.write(srt_output)

    # Delete uploaded video to save space
    os.remove(temp_file.name)

    return srt_file_path

def format_timestamp(seconds: float) -> str:
    """Convert seconds to SRT timestamp format."""
    millisec = int((seconds - int(seconds)) * 1000)
    return f"{int(seconds // 3600):02}:{int((seconds % 3600) // 60):02}:{int(seconds % 60):02},{millisec:03}"

# Gradio interface
demo = gr.Interface(
    fn=transcribe,
    inputs=[
        gr.File(type="binary", label="Upload Video"),
        gr.Dropdown(choices=["tiny", "base", "small", "medium", "large",  "turbo"], value="tiny", label="Choose Whisper Model"),
        gr.Dropdown(
            choices=[
                ("Afrikaans", "af"),
                ("Albanian", "sq"),
                ("Amharic", "am"),
                ("Arabic", "ar"),
                ("Armenian", "hy"),
                ("Azerbaijani", "az"),
                ("Basque", "eu"),
                ("Belarusian", "be"),
                ("Bengali", "bn"),
                ("Bosnian", "bs"),
                ("Bulgarian", "bg"),
                ("Burmesse", "my"),
                ("Catalan", "ca"),
                ("Chinese", "zh"),
                ("Croatian", "hr"),
                ("Czech", "cs"),
                ("Danish", "da"),
                ("Dutch", "nl"),
                ("English", "en"),
                ("Estonian", "et"),
                ("Finnish", "fi"),
                ("French", "fr"),
                ("Galician", "gl"),
                ("Georgian", "ka"),
                ("German", "de"),
                ("Greek", "el"),
                ("Hebrew", "he"),
                ("Hindi", "hi"),
                ("Hungarian", "hu"),
                ("Icelandic", "is"),
                ("Indonesian", "id"),
                ("Irish", "ga"),
                ("Italian", "it"),
                ("Japanese", "ja"),
                ("Kannada", "kn"),
                ("Kazakh", "kk"),
                ("Khmer", "km"),
                ("Korean", "ko"),
                ("Latin", "la"),
                ("Latvian", "lv"),
                ("Lithuanian", "lt"),
                ("Macedonian", "mk"),
                ("Malay", "ms"),
                ("Malayalam", "ml"),
                ("Maori", "mi"),
                ("Marathi", "mr"),
                ("Mongolian", "mn"),
                ("Nepali", "ne"),
                ("Norwegian", "no"),
                ("Persian", "fa"),
                ("Polish", "pl"),
                ("Portuguese", "pt"),
                ("Romanian", "ro"),
                ("Russian", "ru"),
                ("Serbian", "sr"),
                ("Slovak", "sk"),
                ("Slovenian", "sl"),
                ("Spanish", "es"),
                ("Swahili", "sw"),
                ("Swedish", "sv"),
                ("Tagalog", "tl"),
                ("Tamil", "ta"),
                ("Telugu", "te"),
                ("Thai", "th"),
                ("Turkish", "tr"),
                ("Ukrainian", "uk"),
                ("Urdu", "ur"),
                ("Vietnamese", "vi"),
                ("Welsh", "cy"),
            ],
            value="en",
            label="Language Code (ISO)"
        )

    ],
    outputs=gr.File(label="Download SRT File"),
    title="Video → SRT Generator",
    description="Upload a video, select Whisper model and language, then download auto-generated subtitles. Video is deleted after processing."
)

if __name__ == "__main__":
    demo.launch()