import gradio as gr
from pydub import AudioSegment
import edge_tts
import os
import asyncio

# Function to get the length of an audio file in seconds
def get_audio_length(audio_file):
    audio = AudioSegment.from_file(audio_file)
    return audio.duration_seconds

# Function to format time for SRT
def format_time(seconds):
    millis = int((seconds % 1) * 1000)
    seconds = int(seconds)
    hrs = seconds // 3600
    mins = (seconds % 3600) // 60
    secs = seconds % 60
    return f"{hrs:02}:{mins:02}:{secs:02},{millis:03}"

# Function to generate SRT with accurate timing per batch
async def generate_accurate_srt(batch_text, batch_num):
    audio_file = f"batch_{batch_num}_audio.wav"
    
    # Generate the audio using edge-tts
    tts = edge_tts.Communicate(batch_text, "en-US-JennyNeural")
    await tts.save(audio_file)

    # Get the actual length of the audio file
    actual_length = get_audio_length(audio_file)

    # Initialize SRT content
    srt_content = ""
    words = batch_text.split()
    start_time = 0.0
    segment_duration = actual_length / len(words) * 10  # Assuming ~10 words per SRT segment

    # Build SRT content with accurate timing
    for i in range(0, len(words), 10):
        segment_words = words[i:i+10]
        end_time = start_time + segment_duration
        srt_content += f"{i // 10 + 1}\n"
        srt_content += f"{format_time(start_time)} --> {format_time(end_time)}\n"
        srt_content += " ".join(segment_words) + "\n\n"
        start_time = end_time

    return srt_content, audio_file

# Batch processing function for SRT and audio generation
async def batch_process_srt_and_audio(script_text):
    batches = [script_text[i:i+500] for i in range(0, len(script_text), 500)]
    all_srt_content = ""
    combined_audio = AudioSegment.empty()

    for batch_num, batch_text in enumerate(batches):
        srt_content, audio_file = await generate_accurate_srt(batch_text, batch_num)
        all_srt_content += srt_content

        # Append the audio of each batch to the combined audio
        batch_audio = AudioSegment.from_file(audio_file)
        combined_audio += batch_audio

        # Clean up the individual batch audio file
        os.remove(audio_file)

    # Export combined audio and SRT
    combined_audio.export("final_audio.wav", format="wav")
    with open("final_subtitles.srt", "w") as srt_file:
        srt_file.write(all_srt_content)

    return "final_subtitles.srt", "final_audio.wav"

# Gradio interface function
async def process_script(script_text):
    srt_path, audio_path = await batch_process_srt_and_audio(script_text)
    return srt_path, audio_path

# Gradio interface setup
app = gr.Interface(
    fn=process_script,
    inputs=gr.Textbox(label="Enter Script Text", lines=10),
    outputs=[
        gr.File(label="Download SRT File"),
        gr.File(label="Download Audio File")
    ],
    description="Upload your script text, and the app will generate audio and an accurate SRT file for download."
)

app.launch()