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, start_offset):
    audio_file = f"batch_{batch_num}_audio.wav"
    
    # Generate the audio using edge-tts
    tts = edge_tts.Communicate(batch_text, "en-US-AndrewNeural", rate="-25%")
    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()
    segment_duration = actual_length / len(words) * 10  # Adjusted for ~10 words per SRT segment
    start_time = start_offset

    # 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 + (batch_num * 100)}\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, start_time

# 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()
    start_offset = 0.0  # Track cumulative time offset for SRT timing

    for batch_num, batch_text in enumerate(batches):
        srt_content, audio_file, end_offset = await generate_accurate_srt(batch_text, batch_num, start_offset)
        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
        start_offset = end_offset  # Update the start offset for the next batch

        # 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, 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"),
        gr.Audio(label="Play Audio")
    ],
    description="Upload your script text, and the app will generate audio with en-US-AndrewNeural voice (Rate: -25%) and an accurate SRT file for download."
)

app.launch()