import gradio as gr
import os
import ffmpeg
import pysrt
import pandas as pd
import requests
import io
from transformers import MarianMTModel, MarianTokenizer
from gradio_client import Client

# Initialize Gradio Client for Whisper JAX
client = Client(src="sanchit-gandhi/whisper-jax")

def fetch_languages(url):
    response = requests.get(url)
    if response.status_code == 200:
        csv_content = response.content.decode('utf-8')
        df = pd.read_csv(io.StringIO(csv_content), delimiter="|", skiprows=2, header=None).dropna(axis=1, how='all')
        df.columns = ['ISO 639-1', 'ISO 639-2', 'Language Name', 'Native Name']
        df['ISO 639-1'] = df['ISO 639-1'].str.strip()
        language_options = [(row['ISO 639-1'], f"{row['ISO 639-1']} - {row['Language Name']}") for index, row in df.iterrows()]
        return language_options
    else:
        return []

def transcript_audio(audio_file, task, return_timestamps, api_name="/predict_1"):
    prediction = client.predict(audio_file=audio_file, task=task, return_timestamps=return_timestamps, api_name=api_name)
    return prediction['transcription'], prediction['transcription_time_s']

def process_video(input_video, video_language, target_language):
    transcription, _ = transcript_audio(input_video, "transcribe", True)
    srt_path = text_to_srt(transcription)
    translated_srt_path = translate_srt(srt_path, video_language, target_language)
    output_video = add_subtitle_to_video(input_video, translated_srt_path)
    return output_video

language_url = "https://huggingface.co/Lenylvt/LanguageISO/resolve/main/iso.md"
video_language_options = fetch_languages(language_url)

with gr.Blocks() as app:
    with gr.Row():
        input_video = gr.Video(label="Video File")
        video_language = gr.Dropdown(choices=video_language_options, label="Language of the Video")
        target_language = gr.Dropdown(choices=video_language_options, label="Language Translated")
        process_btn = gr.Button("Process Video")
    output_video = gr.Video(label="Video with Translated Subtitles")

    process_btn.click(fn=process_video, inputs=[input_video, video_language, target_language], outputs=output_video)

app.launch()