import os
import yt_dlp
import logging

logging.basicConfig(level=logging.INFO, format='%(levelname)s: %(message)s')

def download_video(url, resolution='HD', output_path='./', cookies_from_browser=None):
    """
    Download a video from a given URL using yt-dlp at a specified resolution.

    Args:
        url (str): The URL of the video to download.
        resolution (str): The desired resolution. Options: '4K', '2K', 'HD'. Default is 'HD'.
        output_path (str): The directory to save the video. Default is current directory.
        cookies_from_browser (str): Browser to extract cookies from (e.g., 'chrome', 'firefox', 'edge').
                                    Use this if you encounter blocking or age-restricted content.

    Returns:
        bool: True if successful, False otherwise.
    """
    # Map resolutions to their approximate vertical pixel height
    resolution_map = {
        '4K': '2160',
        '2K': '1440',
        'HD': '1080'
    }
    
    height = resolution_map.get(resolution.upper(), '1080')
    
    # We ask for the best video with height <= the specified height,
    # plus the best audio, and merge them. If that's not available, 
    # it falls back to the best single file with height <= specified height.
    format_string = f'bestvideo[height<={height}]+bestaudio/best[height<={height}]'
    
    ydl_opts = {
        'format': format_string,
        'outtmpl': os.path.join(output_path, '%(title)s.%(ext)s'),
        'merge_output_format': 'mp4',  # standardizing on mp4 where possible
        'quiet': False,
        'no_warnings': False,
        'nocheckcertificate': True,
        'ignoreerrors': True, # skip unavailable videos in playlists
        'geo_bypass': True,   # try to bypass geographic restrictions
        # Adds sleep to reduce likelihood of getting rate-limited or banned
        'sleep_interval_requests': 1,
        'sleep_interval': 1,
        'max_sleep_interval': 3,
        # Fake user agent can help avoid bot detection
        'http_headers': {
            'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/115.0.0.0 Safari/537.36',
        }
    }
    
    if cookies_from_browser:
        ydl_opts['cookiesfrombrowser'] = (cookies_from_browser, )
    elif os.path.exists('cookies.txt'):
        ydl_opts['cookiefile'] = 'cookies.txt'
        
    try:
        logging.info(f"Starting download for {url} at max resolution {resolution} ({height}p)")
        with yt_dlp.YoutubeDL(ydl_opts) as ydl:
            info = ydl.extract_info(url, download=True)
            if info:
                filename = None
                if 'requested_downloads' in info:
                    filename = info['requested_downloads'][0]['filepath']
                else:
                    filename = ydl.prepare_filename(info)
                    if ydl_opts.get('merge_output_format'):
                        filename = os.path.splitext(filename)[0] + '.' + ydl_opts['merge_output_format']
                
                logging.info(f"Successfully downloaded to {filename}")
                print(os.path.abspath(filename))
                return True
            else:
                logging.error(f"yt-dlp encountered an error.")
                return False
    except yt_dlp.utils.DownloadError as e:
        logging.error(f"Download error: {e}")
        return False
    except Exception as e:
        logging.error(f"Unexpected error: {e}")
        return False

if __name__ == "__main__":
    import argparse
    parser = argparse.ArgumentParser(description="Download videos using yt-dlp.")
    parser.add_argument("url", help="URL of the video to download")
    parser.add_argument("--resolution", "-r", choices=["4K", "2K", "HD"], default="HD", help="Target resolution (4K, 2K, HD). Default: HD")
    parser.add_argument("--output", "-o", default="./", help="Output directory path. Default: current directory")
    parser.add_argument("--cookies", "-c", help="Browser to extract cookies from (e.g. chrome, firefox, edge) to prevent being blocked.")
    
    args = parser.parse_args()
    
    if not os.path.exists(args.output):
        os.makedirs(args.output, exist_ok=True)
        
    download_video(args.url, args.resolution, args.output, args.cookies)