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

# Supported formats
supported_formats = ['avi', 'mp4', 'mov', 'mkv', 'flv', 'wmv', 'webm', 'mpeg', 'mpg', '3gp']
audio_formats = ['mp3', 'wav', 'aac', 'flac']
gif_formats = ['gif']
image_formats = sorted(Image.SAVE.keys()) or ['BLP', 'BMP', 'BUFR', 'DDS', 'DIB', 'EPS',
                                              'GIF', 'GRIB', 'HDF5', 'ICNS', 'ICO', 'IM',
                                              'JPEG', 'JPEG2000', 'MPO', 'MSP', 'PALM', 'PCX',
                                              'PDF', 'PNG', 'PPM', 'SGI', 'SPIDER', 'TGA', 'TIFF',
                                              'WEBP', 'WMX', 'XBM']

def sanitize_filename(filename):
    """Sanitize filename by removing special characters and spaces."""
    filename = re.sub(r'[^a-zA-Z0-9_.-]', '_', filename)  # Replace invalid characters with '_'
    return filename

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')
    duration_in_seconds = float(probe['streams'][0]['duration'])
    return duration_in_seconds

def convert_video(video, target_format, conversion_type, time_in_seconds=None):
    try:
        # Create a temporary directory for the uploaded file
        temp_dir = tempfile.mkdtemp()

        # Sanitize the filename and save the uploaded video to the temp directory
        base_name = os.path.splitext(os.path.basename(video.name))[0]
        sanitized_base_name = sanitize_filename(base_name)
        video_path = os.path.join(temp_dir, f"{sanitized_base_name}.mp4")  # Saving as mp4 by default for now

        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_{sanitized_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_{sanitized_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_{sanitized_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_{sanitized_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_{sanitized_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)

        # Create a temporary file and save uploaded video
        temp_video_path = os.path.join(tempfile.mkdtemp(), video_file.name)
        with open(temp_video_path, "wb") as f:
            f.write(video_file.getbuffer())

        # Get video duration
        video_duration = get_video_duration(temp_video_path)

        # 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.slider(
                label="Time (in seconds) for image extraction",
                min_value=0, 
                max_value=int(video_duration),
                value=0,
                step=1
            )
        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()