"""
A tool for downloading datasets from BigQuery.

Setup:
```
    pip install --upgrade google-cloud-bigquery
    pip install --upgrade google-cloud-storage
    conda install google-cloud-sdk -c conda-forge
```

Note: After installing above packages, run `gcloud auth application-default login` on your terminal.
You will be prompted with a authorization link that you will need to follow and approve using your
email address. Then you will copy-paste authorization code to the terminal.

How to run:
```
python download_bigquery.py <NAME OF YOUR GCP PROJECT> <GCP BIGQUERY DATASET ID> \
    <PATH TO LOCAL FOLDER WHERE DATASET WHERE DATASET WILL BE DOWNLOADED>
    --excluded_tables <(Optional) NAME OF TABLE 1 TO BE IGNORED> <(Optional) NAME OF TABLE 2 TO BE IGNORED>
```

Example: python download_bigquery.py som-nero-nigam-starr \
    som-rit-phi-starr-prod.starr_omop_cdm5_deid_1pcent_lite_2023_02_08 .
"""

from __future__ import annotations

import argparse
import hashlib
import os
import random
import threading
from functools import partial

import google
from google.cloud import bigquery, storage

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Download a Google BigQuery dataset")
    parser.add_argument(
        "gcp_project_name",
        type=str,
        help=(
            "The name of *YOUR* GCP project (e.g. 'som-nero-nigam-starr')."
            " Note that this need NOT be the GCP project that contains the dataset."
            " It just needs to be a GCP project where you have Bucket creation + BigQuery creation permissions."
        ),
    )
    parser.add_argument(
        "gcp_dataset_id",
        type=str,
        help=(
            "The Dataset ID of the GCP dataset to download"
            " (e.g. 'som-rit-phi-starr-prod.starr_omop_cdm5_deid_2022_12_03')."
            " Note that this is the full ID of the dataset (project name + dataset name)"
        ),
    )
    parser.add_argument(
        "output_dir",
        type=str,
        help=(
            "Path to output directory. Note: The downloaded files will be saved in a subdirectory of this,"
            " i.e. `output_dir/gcp_dataset_id/...`"
        ),
    )
    parser.add_argument(
        "--excluded_tables",
        type=str,
        nargs="*",  # 0 or more values expected => creates a list
        default=[],
        help=(
            "Optional. Name(s) of tables to exclude. List tables separated by spaces,"
            " i.e. `--excluded_tables observation note_nlp`"
        ),
    )
    parser.add_argument(
        "--scratch_bucket_postfix",
        type=str,
        default="_extract_scratch",
        help="The postfix for the GCP bucket used for storing temporary files while downloading.",
    )
    args = parser.parse_args()

    target = f"{args.output_dir}/{args.gcp_dataset_id}"
    os.mkdir(target)

    print('Make sure to run "gcloud auth application-default login" before running this command')

    # Connect to our BigQuery project
    client = bigquery.Client(project=args.gcp_project_name)
    storage_client = storage.Client(project=args.gcp_project_name)

    random_dir = hashlib.md5(random.randbytes(16)).hexdigest()

    scratch_bucket_name = args.gcp_project_name.replace("-", "_") + args.scratch_bucket_postfix

    print(f"Storing temporary files in gs://{scratch_bucket_name}/{random_dir}")

    try:
        bucket = storage_client.get_bucket(scratch_bucket_name)
    except google.api_core.exceptions.NotFound as e:
        print(f"Could not find the requested bucket? gs://{scratch_bucket_name} in project {args.gcp_project_name}")
        raise e

    # Get list of all tables in this GCP dataset
    # NOTE: the `HTTPIterator` can be iterated over like a list, but only once (it's a generator)
    tables: google.api_core.page_iterator.HTTPIterator = client.list_tables(args.gcp_dataset_id)
    print(f"Downloading dataset {args.gcp_dataset_id} using your project {args.gcp_project_name}")

    # Use GZIP compression and export as CVSs
    extract_config = bigquery.job.ExtractJobConfig(
        compression=bigquery.job.Compression.GZIP,
        destination_format=bigquery.job.DestinationFormat.CSV,
        field_delimiter=",",
    )

    sem = threading.Semaphore(value=0)  # needed for keeping track of how many tables have been downloaded

    def download(table_id: str, f):
        """Download the results (a set of .csv.gz's) of the BigQuery extract job to our local filesystem
        Note that a single table will be extracted into possibly dozens of smaller .csv.gz files

        Args:
            table_id (str): Name of table (e.g. "attribute_definition")
        """
        if f.errors is not None:
            print("Could not extract, got errors", f.errors, "for", table_id)
            os.abort()
        sem.release()

    n_tables: int = 0
    for table in tables:
        # Get the full name of the table
        table_name: str = f"{table.project}.{table.dataset_id}.{table.table_id}"
        if table.table_id in args.excluded_tables:
            print(f"Skipping extraction | table = {table.table_id}")
            continue
        print(f"Extracting | table = {table.table_id}")
        # Create Google Cloud Storage bucket to extract this table into
        bucket_target_path: str = f"gs://{scratch_bucket_name}/{random_dir}/{table.table_id}/*.csv.gz"
        extract_job = client.extract_table(table.reference, bucket_target_path, job_config=extract_config)
        # Call the `download()` function asynchronously to download the bucket contents to our local filesystem
        extract_job.add_done_callback(partial(download, table.table_id))
        n_tables += 1

    print(f"\n** Extracting a total of {n_tables} tables**\n")
    for i in range(1, n_tables + 1):
        sem.acquire()
        print(f"====> Finished extracting {i} out of {n_tables} tables")

    print("Starting to download tables")

    os.system(f"gsutil -m rsync -r gs://{scratch_bucket_name}/{random_dir} {target}")

    print("------\n------")
    print("Successfully downloaded all tables!")
    print("------\n------")

    # Delete the temporary Google Cloud Storage bucket
    print("\nDeleting temporary files...")
    os.system(f"gsutil -m rm -r gs://{scratch_bucket_name}/{random_dir}")
    print("------\n------")
    print("Successfully deleted temporary Google Cloud Storage files!")
    print("------\n------")