from lxml.html import fromstring
from yta_general_utils.file_processor import replace_file_extension, get_file_extension, write_file_by_chunks_from_response

import requests

def download_file(url: str, output_filename: str):
    """
    Receives a downloadable url as 'url' and downloads that file in
    our system as 'output_filename'.
    """
    if not url:
        return None
    
    if not output_filename:
        return None
    
    r = requests.get(url)

    with open(output_filename, 'wb') as outfile:
        outfile.write(r.content)

    return True

def download_file_from_google_drive(drive_url: str, output_filename: str):
    """
    Downloads a file from a given Google Drive url that must be
    generated by sharing the file and making it public (this
    means anyone can view and download it).

    The file will be stored locally as 'output_filename', but 
    the extension could be changed if the provided one is
    different than the one that the Google Drive file has.
    This final created file name will be returned.

    A valid 'drive_url' must be like this:
    https://drive.google.com/file/d/1rcowE61X8c832ynh0xOt60TH1rJfcJ6z/view?usp=sharing&confirm=1
    """
    # TODO: Check if drive_url is valid and downloadable (regex?)

    # TODO: I think this below is not necessary at all (but it is working hehe)
    if not 'confirm=1' in drive_url:
        drive_url += '&confirm=1'

    # Lets get file id and file name
    drive_file_id = drive_url.replace('https://drive.google.com/file/d/', '').split('/')[0]
    # We fire a simple get request to obtain file name from url title
    drive_file_name = fromstring(requests.get(drive_url).content).findtext('.//title').split('-')[0].strip()

    DOWNLOAD_URL = "https://docs.google.com/uc?export=download&confirm=1"

     # Lets download the file
    session = requests.Session()
    
    # Trying to obtain the web title to get the file name
    response = session.get(DOWNLOAD_URL, params = {'id': drive_file_id}, stream = True)

    # Look for a token to be able to download
    token = None
    for key, value in response.cookies.items():
        if key.startswith('download_warning'):
            token = value
            break

    if token:
        params = {'id': drive_file_id, 'confirm': token}
        response = session.get(DOWNLOAD_URL, params = params, stream = True)

    # We make sure the output extension fits the real one
    output_filename = replace_file_extension(output_filename, get_file_extension(drive_file_name))

    # Save response
    write_file_by_chunks_from_response(response, output_filename)

    return output_filename

def download_audio(url, output_filename):
    """
    Receives a downloadable url as 'url' and downloads that audio in
    our system as 'output_filename'.
    """
    # TODO: Check formats and that it is a valid audio (and use Enums pls)
    if not output_filename.endswith('.mp3') and not output_filename.endswith('.wav') and not output_filename.endswith('.m4a'):
        output_filename = output_filename + '.mp3'

    return download_file(url, output_filename)

def download_image(url: str, output_filename: str):
    """
    Receives a downloadable url as 'url' and downloads that image in
    our system as 'output_filename'.
    """
    if not url:
        return None
    
    if not output_filename:
        return None
    
    # TODO: Check formats and that it is a valid image (and use Enums pls)
    if not output_filename.endswith('.jpg') and not output_filename.endswith('.png') and not output_filename.endswith('.jpeg') and not output_filename.endswith('.webp') and not output_filename.endswith('.bmp'):
        output_filename = output_filename + '.jpg'

    return download_file(url, output_filename)

def download_image_2(url: str, output_filename: str):
    """
    Downloads an image from the provided 'url' if possible.

    # TODO: I need to check this one and compare with the
    # other method so maybe preserve only one (?)
    """
    if not url:
        return None

    if not output_filename:
        return None
    
    import shutil
    
    res = requests.get(url, stream = True)

    if res.status_code == 200:
        with open(output_filename,'wb') as f:
            shutil.copyfileobj(res.raw, f)
        print('Image sucessfully Downloaded: ', output_filename)
    else:
        print('Image Couldn\'t be retrieved')

def download_video(url, output_filename):
    """
    Receives a downloadable url as 'url' and downloads that video in
    our system as 'output_filename'.
    """
    # TODO: Check formats and that it is a valid audio (and use Enums pls)
    if not output_filename.endswith('.mp4') and not output_filename.endswith('.mov'):
        output_filename = output_filename + '.mp4'

    return download_file(url, output_filename)