import whisper
import gradio as gr
import tempfile
import os
from pathlib import Path

MODEL_SIZE = "medium"  # Switched from "small" to "medium"

class GradioSRTGenerator:
    def __init__(self):
        print(f"Loading Whisper model: {MODEL_SIZE}")
        self.model = whisper.load_model(MODEL_SIZE)
        print("Model loaded successfully")

    @staticmethod
    def format_time(seconds: float) -> str:
        """Convert seconds to SRT time format (HH:MM:SS,mmm)"""
        hours = int(seconds // 3600)
        minutes = int((seconds % 3600) // 60)
        secs = int(seconds % 60)
        millis = int((seconds - int(seconds)) * 1000)
        return f"{hours:02d}:{minutes:02d}:{secs:02d},{millis:03d}"

    def build_srt_file(self, segments, srt_path: str):
        """Write Whisper segments into an SRT file at srt_path."""
        with open(srt_path, "w", encoding="utf-8") as srt_file:
            for i, seg in enumerate(segments, start=1):
                start_time = self.format_time(seg["start"])
                end_time = self.format_time(seg["end"])
                text = seg["text"].strip()
                srt_file.write(f"{i}\n")
                srt_file.write(f"{start_time} --> {end_time}\n")
                srt_file.write(f"{text}\n\n")

    def generate_srt(self, video_path, progress=gr.Progress()):
        """
        Transcribe the uploaded video and produce an SRT file.
        Returns: (path_to_srt, status_message)
        """
        if not video_path:
            return None, "⚠️ Please upload a video first."

        try:
            progress(0.1, desc="Transcribing audio with Whisper...")
            result = self.model.transcribe(video_path, task="translate")
            progress(0.6, desc="Building SRT file...")

            # Create a temporary file to hold the SRT
            with tempfile.NamedTemporaryFile(delete=False, suffix=".srt", mode="w", encoding="utf-8") as tmp:
                srt_path = tmp.name

            # Write segments into that SRT
            self.build_srt_file(result["segments"], srt_path)
            progress(1.0, desc="✅ SRT ready!")

            return srt_path, "✅ Transcription complete. Download your SRT below."

        except Exception as e:
            return None, f"❌ Error during transcription: {str(e)}"


def create_ui():
    generator = GradioSRTGenerator()

    with gr.Blocks(theme=gr.themes.Base(primary_hue="blue", secondary_hue="indigo")) as app:
        gr.Markdown(
            """
            # 📝 Video → SRT Generator (Queued)

            Only one transcription job will run at a time; additional users will be placed in a queue.
            Upload any supported video file (MP4, MOV, AVI, MKV, etc.) and click **Generate SRT**. 
            Whisper will transcribe and produce an SRT subtitle file you can download immediately.
            """
        )

        with gr.Row():
            with gr.Column():
                input_video = gr.Video(label="Upload Video")
                generate_btn = gr.Button("🎯 Generate SRT", variant="primary")
                status_text = gr.Textbox(label="Status", interactive=False, show_copy_button=True)

            with gr.Column():
                srt_download = gr.File(label="Download SRT", interactive=False, visible=False)

        # Bind the button and wrap the handler in .queue()
        generate_btn.click(
            fn=generator.generate_srt,
            inputs=[input_video],
            outputs=[srt_download, status_text],
            api_name="generate_srt"
        ).queue()  # ensures this function is queued if another request is running

    # Add a queue to the entire Blocks app with exactly 1 worker
    app.queue(concurrency_count=1, max_size=10)
    return app

app = create_ui()

if __name__ == "__main__":
    app.launch(
        server_name="0.0.0.0",  # bind to all interfaces (important for Docker)
        server_port=7860        # default Gradio port
    )