#!/usr/bin/env python

import os
import time
import argparse
from google.cloud import bigquery, storage
from datetime import date, datetime, timedelta
from gcloud_utils.bigquery.bigquery import Bigquery
from gcloud_utils.storage import Storage

DEFAULT_DATE_FORMAT='%Y%m%d'

os.environ['TZ'] = 'America/Sao_Paulo'
time.tzset()

def main():
    parser = argparse.ArgumentParser(description='Process Bigquery table to Google Cloud Storage')

    parser.add_argument(
        'bigquery_dataset',
        help='Name of BigQuery dataset.'
    )

    parser.add_argument(
        'bigquery_tablename',
        help='Name of BigQuery table.'
    )

    parser.add_argument(
        'cloudstorage_bucket',
        help='Name of CloudStorage bucket.'
    )

    parser.add_argument(
        'cloudstorage_filename',
        help='Name file to store in CloudStorage.'
    )

    parser.add_argument(
        'gcs_key_json',
        help='Name of Json key.'
    )

    parser.add_argument(
        'start_date',
        help='Script start date.',
        default=date.today().strftime(DEFAULT_DATE_FORMAT)
    )

    parser.add_argument(
        'time_delta',
        help='How many days ago you want to get the table.'
    )

    parser.add_argument(
        'export_format',
        help='Export Format.',
        default="csv",
        choices=['csv', 'json', 'avro']

    )

    parser.add_argument(
        'compression_format',
        help='Compression Format.',
        default="gz",
        choices=['gz', 'none', 'snappy']
    )

    args = parser.parse_args()

    dt = datetime.strptime(args.start_date, DEFAULT_DATE_FORMAT) - timedelta(int(args.time_delta))

    seven_days_before = datetime.strptime(args.start_date, DEFAULT_DATE_FORMAT) - timedelta(7)

    seven_days_before_string = seven_days_before.strftime(DEFAULT_DATE_FORMAT)

    start_date_string = dt.strftime(DEFAULT_DATE_FORMAT)

    bq_table_name = args.bigquery_tablename % dict(date=start_date_string)

    bq_client = Bigquery(bigquery.Client.from_service_account_json(args.gcs_key_json))

    compression_format = None if args.compression_format=="none" else args.compression_format

    storage_client = Storage(args.cloudstorage_bucket, storage.Client.from_service_account_json(args.gcs_key_json))

    storage_client.rename_files(args.cloudstorage_filename, "tmp_{}/{}".format(start_date_string, args.cloudstorage_filename))
    

    bq_client.table_to_cloud_storage(
        args.bigquery_dataset, bq_table_name, args.cloudstorage_bucket, args.cloudstorage_filename,
        export_format=args.export_format,
        compression_format=compression_format
    )

    try:
        storage_client.delete_path("tmp_{}/{}".format(seven_days_before_string, args.cloudstorage_filename))
    except:
        pass

if __name__ == "__main__":
    main()
