import streamlit as st
import ffmpeg
import os
from PIL import Image

# Supported formats
supported_formats = ['avi', 'mp4', 'mov', 'mkv', 'flv', 'wmv', 'webm', 'mpeg', 'mpg', '3gp']
audio_formats = ['mp3', 'wav', 'aac', 'flac']
image_formats = sorted(Image.SAVE.keys())  # Supported image formats

def get_video_duration(video_path):
    """Get video duration in seconds using ffmpeg."""
    probe = ffmpeg.probe(video_path, v='error', select_streams='v:0', show_entries='stream=duration')
    return float(probe['streams'][0]['duration'])

def convert_video(video, target_format, conversion_type, time_in_seconds=None):
    try:
        base_name = os.path.splitext(os.path.basename(video.name))[0]
        video_path = f"./temp/{base_name}"

        with open(video_path, "wb") as f:
            f.write(video.getbuffer())  # Save the uploaded video to a local file

        if conversion_type == 'Video to Video':
            output_file = f"flowly_ai_video_converter_{base_name}.{target_format.lower()}"
            ffmpeg.input(video_path).output(output_file).run()  # Use ffmpeg.input properly
            return output_file

        elif conversion_type == 'Video to Audio':
            audio_output_file = f"flowly_ai_video_to_audio_{base_name}.{target_format.lower()}"
            ffmpeg.input(video_path).output(audio_output_file).run()  # Use ffmpeg.input properly
            return audio_output_file

        elif conversion_type == 'Video to GIF':
            gif_output_file = f"flowly_ai_video_to_gif_{base_name}.gif"
            ffmpeg.input(video_path).output(gif_output_file, vf="fps=10,scale=320:-1:flags=lanczos").run()  # Use ffmpeg.input properly
            return gif_output_file

        elif conversion_type == 'Video to Image':
            if time_in_seconds is None:
                return "Please specify a valid time in seconds for image extraction."

            image_output_file = f"flowly_ai_video_to_image_{base_name}_{time_in_seconds}.png"
            ffmpeg.input(video_path, ss=time_in_seconds).output(image_output_file, vframes=1).run()  # Use ffmpeg.input properly

            # Convert the image to the desired format using Pillow
            img = Image.open(image_output_file)
            if target_format.lower() in image_formats:
                image_output_file = f"flowly_ai_video_to_image_{base_name}_{time_in_seconds}.{target_format.lower()}"
                img.save(image_output_file, target_format.upper())  # Save with the desired format

            return image_output_file

    except Exception as e:
        return f"Error: {e}"

def update_format_choices(conversion_type):
    """Update available format choices based on the conversion type."""
    if conversion_type == 'Video to Video':
        return supported_formats
    elif conversion_type == 'Video to Audio':
        return audio_formats
    elif conversion_type == 'Video to GIF':
        return gif_formats
    elif conversion_type == 'Video to Image':
        return image_formats
    return []

def main():
    st.title("Video Conversion Tool")

    # Upload video file
    video_file = st.file_uploader("Upload a Video", type=supported_formats)

    if video_file:
        st.video(video_file)

        # Select conversion type
        conversion_type = st.selectbox(
            "Select Conversion Type", 
            ['Video to Video', 'Video to Audio', 'Video to GIF', 'Video to Image']
        )

        # Update format choices based on conversion type
        target_format_choices = update_format_choices(conversion_type)
        target_format = st.selectbox("Select Target Format", target_format_choices)

        # If 'Video to Image' conversion, ask for time in seconds
        if conversion_type == 'Video to Image':
            time_in_seconds = st.number_input("Time (in seconds) for image extraction", min_value=0)
        else:
            time_in_seconds = None

        if st.button("Convert"):
            with st.spinner("Converting..."):
                output_file = convert_video(video_file, target_format, conversion_type, time_in_seconds)
                
                if "Error" in output_file:
                    st.error(output_file)
                else:
                    st.success(f"Conversion Successful! Download the file:")
                    st.download_button(label="Download Converted File", data=open(output_file, 'rb').read(), file_name=output_file)

if __name__ == "__main__":
    main()