import streamlit as st
from pytube import Playlist, YouTube
import time
import os
import re

def sanitize_filename(filename):
    # Remove invalid characters and replace spaces with underscores
    return re.sub(r'[\\/*?:"<>|]', "", filename.replace(" ", "_"))

def download_video(URL, save_folder):
    try:
        video = YouTube(URL)
        video_title = sanitize_filename(video.title)
        video_filename = f"{video_title}.mp4"
        file_path = os.path.join(save_folder, video_filename)
        
        st.write(f"Checking if file {file_path} exists...")
        if os.path.exists(file_path):
            st.write("File exists. Skipping download.")
            return
        else:
            st.write("File does not exist. Proceeding with download.")

        st.write(f"\nDownloading your File: {video_title}")

        if not os.path.exists(save_folder):
            os.makedirs(save_folder)

        st.write("Downloading video...")
        video.streams.get_highest_resolution().download(output_path=save_folder, filename=video_title)
        st.write('Hurray! Successfully Downloaded')
    except Exception as e:
        st.write("Error occurred while downloading the video:", str(e))

    return True


def download_playlist(playlist_url, save_folder):
    try:
        playlist = Playlist(playlist_url)
        st.write(f'Downloading your Playlist: {playlist.title}')

        start_time = time.time()
        total_videos = len(playlist.video_urls)
        downloaded_videos = 0

        for video_url in playlist.video_urls:
            video = YouTube(video_url)
            video_title = sanitize_filename(video.title)
            video_filename = f"{video_title}.mp4"
            file_path = os.path.join(save_folder, video_filename)

            st.write(f"Checking if file {file_path} exists...")
            if os.path.exists(file_path):
                st.write(f"Video '{video_title}' already exists, skipping...")
                downloaded_videos += 1
                continue
            else:
                st.write("File does not exist. Proceeding with download.")

            st.write(f"Downloading video: {video_title}")
            video.streams.get_highest_resolution().download(output_path=save_folder, filename=video_title)
            downloaded_videos += 1
            st.write(f"Downloaded video {downloaded_videos}/{total_videos}")

        end_time = time.time()

        st.write(f'Hurray! Successfully Downloaded Playlist: {playlist.title}')
        st.write(f'Total Time Taken: {end_time - start_time:.2f} seconds')
    except Exception as e:
        st.write("Error occurred while processing the playlist:", str(e))

    return True

def main():
    st.title("YouTube Downloader")
    download_type = st.selectbox("What do you want to download?", ["Single Video", "Playlist"])

    if download_type == "Single Video":
        URL = st.text_input("Enter the YouTube video URL")
        save_folder = st.text_input("Enter the folder path where you want to save the video")
        if st.button("Download Video"):
            if download_video(URL, save_folder):
                st.write("Do you want to download another video?")
    else:
        playlist_url = st.text_input("Enter the YouTube playlist URL")
        save_folder = st.text_input("Enter the folder path where you want to save the videos")
        if st.button("Download Playlist"):
            if download_playlist(playlist_url, save_folder):
                st.write("Do you want to download another playlist?")

if __name__ == "__main__":
    main()