import streamlit as st
import tempfile
import os
import ffmpeg  # Make sure to install ffmpeg-python

def save_uploaded_file(uploaded_file):
    try:
        with tempfile.NamedTemporaryFile(delete=False, suffix=os.path.splitext(uploaded_file.name)[1]) as tmp_file:
            tmp_file.write(uploaded_file.read())
            return tmp_file.name
    except Exception as e:
        st.error(f"Error saving uploaded file: {e}")
        return None

def add_subtitle_to_video(input_video_path, subtitle_file_path, subtitle_language, soft_subtitle):
    output_video_path = os.path.join(tempfile.gettempdir(), f"output-{os.path.basename(input_video_path)}")
    try:
        if soft_subtitle:
            # For soft subtitles, use the mov_text codec and add metadata
            stream = ffmpeg.input(input_video_path).output(
                ffmpeg.input(subtitle_file_path),
                output_video_path,
                **{'c': 'copy', 'c:s': 'mov_text'},
                **{'metadata:s:s:0': f'language={subtitle_language}'}
            )
        else:
            # For hard subtitles, burn them into the video stream
            stream = ffmpeg.input(input_video_path).output(
                output_video_path,
                **{'vf': f'subtitles={subtitle_file_path}'}
            )
        ffmpeg.run(stream, overwrite_output=True)
    except ffmpeg.Error as e:
        st.error(f"FFmpeg error: {e.stderr}")
        return None
    return output_video_path

def process_video(video_file, subtitle_file, subtitle_type, subtitle_language):
    if video_file is not None and subtitle_file is not None:
        video_path = save_uploaded_file(video_file)
        subtitle_path = save_uploaded_file(subtitle_file)
        if video_path and subtitle_path:
            soft_subtitle = subtitle_type == "Soft"
            processed_video_path = add_subtitle_to_video(video_path, subtitle_path, subtitle_language, soft_subtitle)
            return processed_video_path
    return None

# Streamlit UI
st.title("Video Subtitler")
st.markdown("The process can be long. For API use, please visit [this space](https://huggingface.co/spaces/Lenylvt/SRT_to_Video-API).")

video_file = st.file_uploader("📹 Upload a video", type=["mp4", "avi", "mov"])
subtitle_file = st.file_uploader("📜 Upload subtitle file", type=["srt", "vtt"])
subtitle_type = st.radio("🧷 Subtitle Type", ["Hard (Directly on video)", "Soft"], index=0)  # Default to "Hard"
subtitle_language = st.text_input("🗯️ Subtitle Language (ISO 639-1 code)", value="en")

if st.button("📁 Process Video"):
    processed_video_path = process_video(video_file, subtitle_file, subtitle_type, subtitle_language)
    if processed_video_path:
        st.video(processed_video_path)
    else:
        st.error("🔴 Please upload both a video and a subtitle file.")