import streamlit as st
import os
import subprocess
from pathlib import Path

# Directories for input and output
INPUT_DIR = 'input_videos'
OUTPUT_DIR = 'output_videos'

# Ensure directories exist
os.makedirs(INPUT_DIR, exist_ok=True)
os.makedirs(OUTPUT_DIR, exist_ok=True)


# Function to clear directories
def clear_directory(directory):
    for file in Path(directory).glob("*"):
        file.unlink()


# Streamlit application
st.title('Video Super Resolution Enhancement')

# Video upload
uploaded_file = st.file_uploader("Upload a video file", type=["mp4", "mov", "avi"])

if uploaded_file:
    input_video_path = os.path.join(INPUT_DIR, uploaded_file.name)

    # Save uploaded video
    with open(input_video_path, "wb") as f:
        f.write(uploaded_file.getbuffer())

    st.video(input_video_path)

    # Run the enhancement command
    if st.button("Enhance Video"):
        command = [
            "python", "inference_realesrgan_video.py",
            "-n", "RealESRGAN_x4plus",
            "-i", input_video_path,
            "--outscale", "2",
            "--face_enhance",
            "-o", OUTPUT_DIR
        ]

        # Run the command and wait for it to complete
        process = subprocess.run(command, capture_output=True, text=True)

        if process.returncode == 0:
            st.success("Video enhanced successfully!")

            # Check for the output video file
            output_files = list(Path(OUTPUT_DIR).glob("*.mp4"))
            if output_files:
                output_video_path = str(output_files[0])
                st.video(output_video_path)

                with open(output_video_path, "rb") as file:
                    btn = st.download_button(
                        label="Download enhanced video",
                        data=file,
                        file_name=Path(output_video_path).name,
                        mime="video/mp4"
                    )

                if btn:
                    # Clear input and output directories after download
                    clear_directory(INPUT_DIR)
                    clear_directory(OUTPUT_DIR)
            else:
                st.error("No output video found in the output directory.")
        else:
            st.error(f"Error enhancing video: {process.stderr}")