import gradio as gr
import os
from moviepy.editor import VideoFileClip, TextClip, CompositeVideoClip
import ffmpeg  # Make sure to install ffmpeg-python

def read_subtitle_file(subtitle_path):
    with open(subtitle_path, 'r', encoding='utf-8') as file:
        subtitle_content = file.read()
    return os.path.basename(subtitle_path), subtitle_content

def add_subtitle_to_video(input_video, subtitle_file, subtitle_language, soft_subtitle):
    video_input_stream = ffmpeg.input(input_video)
    subtitle_input_stream = ffmpeg.input(subtitle_file)
    input_video_name = os.path.splitext(os.path.basename(input_video))[0]
    output_video = f"/tmp/output-{input_video_name}.mp4"
    subtitle_track_title = os.path.splitext(os.path.basename(subtitle_file))[0]

    if soft_subtitle:
        stream = ffmpeg.output(
            video_input_stream, subtitle_input_stream, output_video,
            **{"c": "copy", "c:s": "mov_text"},
            **{"metadata:s:s:0": f"language={subtitle_language}",
               "metadata:s:s:0": f"title={subtitle_track_title}"}
        )
    else:
        stream = ffmpeg.output(
            video_input_stream, output_video,
            vf=f"subtitles={subtitle_file}"
        )

    ffmpeg.run(stream, overwrite_output=True)
    return output_video

def video_demo(video, subtitle, subtitle_type, subtitle_language):
    if subtitle is not None:
        soft_subtitle = subtitle_type == "Soft"
        processed_video_path = add_subtitle_to_video(video, subtitle, subtitle_language, soft_subtitle)
        return processed_video_path
    else:
        return video

demo = gr.Interface(
    fn=video_demo,
    inputs=[
        gr.Video(label="Video", interactive=True),
        gr.File(label="Subtitle", file_types=[".srt", ".vtt"]),
        gr.Radio(choices=["Soft", "Hard"], label="Subtitle Type"),
        gr.Textbox(label="Subtitle Language (ISO 639-1 code, e.g., 'en' for English)", default="en"),
    ],
    outputs=gr.Video(label="Processed Video"),
)

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