import requests
from tqdm import tqdm
import zipfile
import os
import shutil

def _create_if_not_exists(path: str, remove=True) -> None:
    if path is None:
        return
        
    if os.path.exists(path):
        if not remove:
            return
        else:
            if os.path.isdir(path):
                shutil.rmtree(path)
            else:
                os.remove(path)

    # If is a file -> get parent. If no parent -> exit
    if os.path.split(path)[0] == '':
        return

    if '.' in os.path.split(path)[1]:
        path = os.path.split(path)[0]

    os.makedirs(path, exist_ok=True)

def download(url: str, filename: str, unzip=False, unzip_path: str = None, force_download=False, force_unzip=False, clean=False) -> None:
    """
    Download a file from a OneDrive url.

    Yes, this documentation was generated by copilot. Thank you copilot.
    :param url: The url to download from (should end with '?download=1').
    :param filename: The filename to save the file as.
    :param unzip: Whether to unzip the file.
    :param unzip_path: The path to unzip the file to. Default is current path.
    :param force_download: Whether to force download the file even if it already exists.
    :param force_unzip: Whether to force unzip and overwrite the file even if it already exists.
    :param clean: Whether to clean the unzipped files after unzipping.
    """
    assert url is not None, "URL cannot be None!"
    assert filename is not None, "Parameter filename cannot be None!"

    if not url.endswith("?download=1"):
        # replace everithing after the last ? with ?download=1
        url = url.split("?")[0] + "?download=1"

    try:
        response = requests.get(url, stream=True)
        if 'id=' in response.url and '&' in response.url:
            fname = response.url.split("id=")[-1].split("&")[0].split("%2F")[-1]
        else:
            fname = response.url.split('/')[-1]

        if os.path.split(filename)[-1] == '' or '.' not in os.path.split(filename)[-1]:
            filename = os.path.join(filename, fname)

        total_size_in_bytes = int(response.headers.get('content-length', 0))
        progress_bar = tqdm(total=total_size_in_bytes, unit='iB', unit_scale=True) if total_size_in_bytes > 1024 else None
        block_size = 1024

        if not os.path.exists(filename) or force_download:
            _create_if_not_exists(filename)

            with open(filename, 'wb') as f:
                for data in response.iter_content(block_size):
                    if progress_bar is not None:
                        progress_bar.update(len(data))
                    f.write(data)
            
            if progress_bar is not None:
                progress_bar.close()

                if total_size_in_bytes != 0 and progress_bar.n != total_size_in_bytes:
                    raise Exception(
                        f"ERROR, something went wrong during download!\nExpected {total_size_in_bytes} Bytes, got {progress_bar.n} Bytes.")

        # unzip file if necessary
        if unzip:
            if filename.endswith(".zip"):
                unzip_path = unzip_path if unzip_path is not None else os.path.split(filename)[0]
                clean_unzip_path = force_unzip and os.path.realpath(unzip_path) not in os.path.realpath(filename)
        
                _create_if_not_exists(unzip_path, remove=clean_unzip_path)

                if force_unzip:
                    print("Warning: overwriting existing files!")

                with zipfile.ZipFile(filename, 'r') as zip_ref:
                    for file in tqdm(iterable=zip_ref.namelist(), total=len(zip_ref.namelist()), desc="Extracting files"):
                        if not os.path.exists(os.path.join(unzip_path, file)) or force_unzip:
                            zip_ref.extract(member=file, path=unzip_path)

                if clean:
                    os.remove(filename)

    except Exception as e:
        print(e)
        raise Exception("ERROR, something went wrong, see error above!")


if __name__ == "__main__":
    ln = "https://unimore365-my.sharepoint.com/:u:/g/personal/215580_unimore_it/EQ-DxzGOF7lBt90A601kvVEBR_ca9PtUdN_asesZ-F80bw?download=1"
    print('Downloading dataset')
    download(ln, filename="./tmp/", unzip=True)
