from pytube import YouTube
import os

def download_video(url, resolution='highest', output_path='downloads'):
    yt = YouTube(url)
    if resolution == 'highest':
        stream = yt.streams.get_highest_resolution()
    else:
        stream = yt.streams.filter(res=resolution).first()
    
    if stream:
        print(f'Downloading {stream.title}...')
        stream.download(output_path=output_path)
        print('Download complete!')
    else:
        print('Resolution not available. Downloading highest resolution...')
        stream = yt.streams.get_highest_resolution()
        stream.download(output_path=output_path)
        print('Download complete!')

def download_audio(url, output_path='downloads'):
    yt = YouTube(url)
    stream = yt.streams.filter(only_audio=True).first()
    
    if stream:
        print(f'Downloading {stream.title}...')
        stream.download(output_path=output_path)
        print('Download complete!')
    else:
        print('Audio stream not available.')

def main():
    os.makedirs('downloads', exist_ok=True)
    url = input('Enter the YouTube URL: ')
    choice = input('Enter "video" to download video or "audio" to download audio: ').lower()

    if choice == 'video':
        resolution = input('Enter resolution (e.g., 720p) or "highest" for best quality: ').lower()
        download_video(url, resolution)
    elif choice == 'audio':
        download_audio(url)
    else:
        print('Invalid choice. Please enter "video" or "audio".')

if __name__ == '__main__':
    main()