import streamlit as st
import os
import yt_dlp
import psutil

# Function to get disk usage information
def get_disk_info():
    disk_info = psutil.disk_usage('/')
    return disk_info

# Streamlit UI setup
st.title("YouTube Video Downloader 8K supported")

# Placeholder for disk usage info
disk_info_placeholder = st.empty()

# Function to update disk usage information
def update_disk_info():
    disk_info = get_disk_info()
    disk_used_gb = disk_info.used / (1024 ** 3)
    disk_total_gb = disk_info.total / (1024 ** 3)
    
    disk_info_placeholder.markdown("## Server Information")
    disk_info_placeholder.markdown(f"**Disk Usage:** {disk_info.percent}% used ({disk_used_gb:.2f} GB of {disk_total_gb:.2f} GB)")

# Initial disk information update
update_disk_info()

# Streamlit input fields
url = st.text_input("Enter the YouTube URL:")
quality = st.selectbox("Select the desired quality:", ['4320', '2160', '1440', '1080', '720', '480'])
save_file_name = st.text_input("Enter the save file name (without extension):")

# Convert quality to the appropriate format string for yt-dlp
quality_format = f"bestvideo[height<={quality}]+bestaudio/best[height<={quality}]"

# Define download options based on user input
ydl_opts = {
    'format': quality_format,  # Download the specified quality
    'outtmpl': f'{save_file_name}.%(ext)s',  # Save the video with the specified name
}

if st.button("Download"):
    with st.spinner("Downloading... Please wait."):
        try:
            # Download the video
            with yt_dlp.YoutubeDL(ydl_opts) as ydl:
                result = ydl.download([url])
            
            # Check if download was successful
            if result == 0:
                # Get the actual file extension used
                info_dict = ydl.extract_info(url, download=False)
                extension = info_dict.get('ext', 'webm')
                file_path = f'{save_file_name}.{extension}'

                # Check if the file exists
                if os.path.exists(file_path):
                    # Provide a download link for the file
                    with open(file_path, "rb") as file:
                        st.download_button(
                            label="Download video",
                            data=file,
                            file_name=f"{save_file_name}.{extension}",
                            mime=f"video/{extension}"
                        )
                    st.success('Your video is ready for download.')
                else:
                    st.error("File not found. Please make sure the download was successful.")
            else:
                st.error("Download failed. Please check the URL and try again.")
        except Exception as e:
            st.error(f"An error occurred: {str(e)}")

# Update disk information whenever the user interacts with the app
update_disk_info()