#!/usr/bin/env python

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

DEFAULT_DATE_FORMAT = '%Y%m%d'

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


def main():
    parser = argparse.ArgumentParser(description='Process Google Cloud Storage to 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(
        'bigquery_dataset',
        help='Name of BigQuery dataset.'
    )

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

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

    parser.add_argument(
        'start_date',
        help='Script start date to use in tablename if you use date in name.',
        default=date.today().strftime(DEFAULT_DATE_FORMAT)
    )

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

    )
    args = parser.parse_args()

    dt = datetime.strptime(args.start_date, DEFAULT_DATE_FORMAT)

    start_date_string = dt.strftime(DEFAULT_DATE_FORMAT)

    bq_table_name = args.bigquery_tablename.format(date=start_date_string)

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

    job_config = bigquery.LoadJobConfig()
    job_config.write_disposition = bigquery.WriteDisposition.WRITE_TRUNCATE
    job_config.autodetect = True

    bq_client.cloud_storage_to_table(
        args.cloudstorage_bucket, args.cloudstorage_filename, args.bigquery_dataset, bq_table_name, job_config, args.import_format
    )


if __name__ == "__main__":
    main()
