#!/usr/bin/env python
"""Download datasets"""
import zipfile
from pathlib import Path
from typing import Literal

import dcargs
import gdown


def download_friends():
    """Script ported from: https://github.com/ethanweber/sitcoms3D/blob/master/download_data.py"""
    data_filename_to_file_id = [
        ("./data/human_data.zip", "1zk12qxnZbReKqdvy-JYp1jQn2QA3oJzz"),
        ("./data/human_pairs.zip", "1927rnmJ3mWcjsj0afqawSto2IK2lLwmA"),
        ("./data/sparse_reconstruction_and_nerf_data.zip", "1RmwDUAp1T4RkwZg1S_7L9JAViWtBZfAG"),
    ]

    disk_data_filename_to_file_id = [  # pylint: disable=unused-variable
        ("./data/sparse_reconstruction_and_nerf_data/ELR-apartment-disk.zip", "1YcTwMv5PP0uqXYdN-Lp8k0WnXUdHBXrY"),
        ("./data/sparse_reconstruction_and_nerf_data/Frasier-apartment-disk.zip", "1vPqplDv5rrFu5LrxTHbET0FiePTGi8zE"),
        (
            "./data/sparse_reconstruction_and_nerf_data/Friends-monica_apartment-disk.zip",
            "1yNpU4M44gIuWEaC5tSAOKy-A6KX8wTm8",
        ),
        (
            "./data/sparse_reconstruction_and_nerf_data/HIMYM-red_apartment-disk.zip",
            "1m76i46YYlpDKvlx-Jn7qDDAvkOVVVmqF",
        ),
        (
            "./data/sparse_reconstruction_and_nerf_data/Seinfeld-jerry_living_room-disk.zip",
            "1YFoV3Gd9asKsRarZvzRTYtri1755rBPx",
        ),
        ("./data/sparse_reconstruction_and_nerf_data/TAAHM-kitchen-disk.zip", "1QpF8kVGfgkqm_cDY5sY_E82BT8VOWL38"),
        (
            "./data/sparse_reconstruction_and_nerf_data/TBBT-big_living_room-disk.zip",
            "175BjkpxMAOt75ZVIjYpArxWpO8uRexgJ",
        ),
    ]

    # !!! Large files !!! Only uncomment if you need the disk correspondences.
    # data_filename_to_file_id += disk_data_filename_to_file_id

    for data_filename, file_id in data_filename_to_file_id:
        # Download the files
        url = f"https://drive.google.com/uc?id={file_id}"
        gdown.download(url, output=data_filename, quiet=False)
        with zipfile.ZipFile(data_filename, "r") as zip_ref:
            zip_ref.extractall(data_filename.replace(".zip", "").replace("-disk", ""))
        Path(data_filename).unlink(missing_ok=True)


def download_blender():
    """Download blender dataset and format accordingly"""
    url = "https://drive.google.com/uc?id=18JxhpWD-4ZmuFKLzKlAw-w5PpzZxXOcG"
    download_path = Path("data/blender_data.zip")
    gdown.download(url, output=str(download_path), quiet=False)
    with zipfile.ZipFile(download_path, "r") as zip_ref:
        zip_ref.extractall("data/")
    unzip_path = Path("data/nerf_synthetic")
    final_path = Path("data/blender")
    unzip_path.rename(final_path)
    download_path.unlink(missing_ok=True)


def main(dataset: Literal["blender", "friends"]):
    """Main download script to download all data"""
    save_dir = Path("data/")
    save_dir.mkdir(parents=True, exist_ok=True)

    if dataset == "blender":
        download_blender()
    if dataset == "friends":
        download_friends()


if __name__ == "__main__":
    try:
        # Requires dcargs>=0.3.3.
        dcargs.extras.set_accent_color("bright_yellow")
    except AttributeError:
        pass

    dcargs.cli(main)