Main

import os
import threading
import moviepy.editor as mp
from pytube import YouTube, Playlist
from tkinter import filedialog
from tkinter import *
from tkinter import ttk
import shelve


def clean_title(title):
    check = ["~", "*", "<", ">", "?", ":", "|", "/", "\\", ".", ",", "\"", "\'"]

    for _ in check:
        if _ in title:
            title = title.replace(_, "")

    return title


def set_download_path():
    new_path = str(filedialog.askdirectory())
    if "/" in new_path:
        new_path = new_path.replace("/", "\\")

    config = shelve.open(".\\config\\config")
    config["download_path"] = new_path

    config.close()


def callback(stream, chunk, file_handle, bytes_remaining):

    global _bytes_remaining
    global filesize

    if int(_bytes_remaining) == 0:
        filesize = bytes_remaining
        _bytes_remaining = filesize
    else:
        _bytes_remaining = bytes_remaining

    progress["maximum"] = filesize
    pval = filesize - _bytes_remaining
    progress["value"] = int(pval)
    progress.update_idletasks()

    # print("Bytes remaining: " + str(bytes_remaining))


def download():
    def inner_download_audio():

        yt = YouTube(url)
        yt.register_on_progress_callback(callback)

        if "&list=" not in url:
            c_title = clean_title(yt.title)

            yt.streams.first().download(".\\temp")

            if os.path.exists(".\\temp\\" + c_title + ".mp4"):
                with mp.VideoFileClip(".\\temp\\" + c_title + ".mp4") as clip:
                    clip.audio.write_audiofile(download_path + "\\" + c_title + ".mp3")
            else:
                yt.streams.filter(only_audio=True).first().download(download_path)

            for file in os.listdir(".\\temp\\"):
                if os.path.splitext(file)[0] == c_title:
                    os.remove(".\\temp\\" + file)
        else:
            pl = Playlist(url)
            pl.download_all(".\\temp\\")

            for file in os.listdir(".\\temp\\"):
                c_title = clean_title(os.path.splitext(file)[0])

                if os.path.exists(".\\temp\\" + c_title + ".mp4"):
                    with mp.VideoFileClip(".\\temp\\" + c_title + ".mp4") as clip:
                        clip.audio.write_audiofile(download_path + "\\" + c_title + ".mp3")
                    os.remove(".\\temp\\" + file)
                else:
                    print("File not found: " + file)

        print("Download Complete")

    def inner_download_video():
        yt = YouTube(url)
        yt.register_on_progress_callback(callback)

        if "&list=" not in url:
            yt.streams.first().download(download_path)
        else:
            pl = Playlist(url)
            pl.download_all(download_path)

        print("Download Complete")

    url = url_entry.get()
    url_entry.delete(0, END)

    config = shelve.open(".\\config\\config")
    download_path = config.get("download_path", "")

    state = url_checkbox_state.get()

    if state == 0:
        tv = threading.Thread(target=inner_download_video)
        tv.start()
    else:
        ta = threading.Thread(target=inner_download_audio)
        ta.start()

    config.close()

_bytes_remaining = 0
filesize = 9

# config
if os.path.exists(".\\config\\config.dir"):
    pass
else:
    config = shelve.open(".\\config\\config")
    config["download_path"] = "."
    config.close()


# tkinter
root = Tk()
root.title("YDownloader")
root.geometry("250x100+650+300")

menu = Menu(root)
root.config(menu=menu)

subMenu = Menu(menu)
menu.add_cascade(label="file", menu=subMenu)
subMenu.add_command(label="change download location", command=set_download_path)
subMenu.add_separator()
subMenu.add_command(label="exit", command=root.destroy)

url_entry = Entry(root, width=50)

url_checkbox_state = IntVar()
url_checkbox = Checkbutton(root, text="Audio only", variable=url_checkbox_state)

download_btn = Button(root, bg="green", text="Download", command=download)

progress = ttk.Progressbar(root, orient=HORIZONTAL, length=225, mode='determinate')

url_entry.pack()
url_checkbox.pack()
download_btn.pack()
progress.pack()

if __name__ == '__main__':

    root.mainloop()

# TODO 1: Add progress bar for playlists