"""download files in bulk with human readable names."""

import argparse
import asyncio
import os
from pathlib import Path

import aioboto3

import snowflake.connector
from snowflake.connector import DictCursor

from src.query import GET_FILES


async def main():
    """Entrypoint."""
    args = get_args()

    # setup snowflake connection
    snowflake_user = os.environ['SNOWFLAKE_USER']
    snowflake_db = os.environ['SNOWFLAKE_DB']
    snowflake_warehouse = os.environ['SNOWFLAKE_WAREHOUSE']
    snowflake_auth_type = os.environ['SNOWFLAKE_AUTH_TYPE']
    snowflake_private_key_file = \
        os.environ.get('SNOWFLAKE_PRIVATE_KEY_FILE')
    snowflake_private_key_file_pwd = \
        os.environ.get('SNOWFLAKE_PRIVATE_KEY_FILE_PWD')

    snowflake_account = 'delphi.us-east-1'
    if snowflake_auth_type not in ('externalbrowser', 'snowflake'):
        print(f'Invalid "SNOWFLAKE_AUTH_TYPE" of "{snowflake_auth_type}"')  # noqa:E501
        exit(1)

    conn_args = {
        'user': snowflake_user,
        'authenticator': snowflake_auth_type,
        'account': snowflake_account
    }
    if snowflake_auth_type == 'snowflake':
        conn_args['private_key_file'] = snowflake_private_key_file
        if snowflake_private_key_file_pwd:
            conn_args['private_key_file_pwd'] = snowflake_private_key_file_pwd
    snowflake_conn = snowflake.connector.connect(**conn_args)
    cursor = snowflake_conn.cursor(DictCursor)
    cursor.execute(f'USE DATABASE {snowflake_db}')
    cursor.execute(f'USE WAREHOUSE {snowflake_warehouse}')

    # handle args
    num_workers = args.workers
    upcs = args.upc_input
    asset_types = args.types

    # query snowflake for data
    cursor.execute(
        GET_FILES,
        {
            'asset_types': asset_types,
            'upcs': upcs
        }
    )
    rows = cursor.fetchall()
    cursor.close()
    snowflake_conn.close()
    print(f'Query returned {len(rows)} rows')

    # queue up data for processing
    queue = asyncio.Queue()
    for row in rows:
        queue.put_nowait((
            aioboto3.Session(),
            row
        ))

    # download data
    await asyncio.gather(
        *[
            asyncio.create_task(process_queue(queue))
            for _ in range(num_workers)
        ]
    )


def get_args():
    """Read and validate CLI args."""
    parser = argparse.ArgumentParser(
        formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    parser.add_argument(
        'upc_input', type=str,
        help='UPC comma separated string or filename containing UPCs (one per line) in ./input')  # noqa:E501
    parser.add_argument(
        '-w', dest='types', action='append_const', const='WAV',
        help='Download mezzanine WAV assets')
    parser.add_argument(
        '-f', dest='types', action='append_const', const='FLAC',
        help='Download mezzanine FLAC assets')
    parser.add_argument(
        '-t', dest='types', action='append_const', const='TIF',
        help='Download mezzanine TIF assets')
    parser.add_argument(
        '-m', dest='types', action='append_const', const='MP3_192',
        help='Download mezzanine MP3 assets')
    parser.add_argument(
        '-j', dest='types', action='append_const', const='JPG',
        help='Download mezzanine JPG assets')
    parser.add_argument(
        '--workers', type=int, default=20,
        help='Number of parallel workers to download with')
    args = parser.parse_args()

    # read in UPCs from file if exists
    upc_filename = f'/tmp/input/{args.upc_input}'
    if os.path.isfile(upc_filename):
        with open(upc_filename, 'r') as f:
            args.upc_input = list(set([
                x.strip()
                for x in f.readlines()
            ]))
    else:
        args.upc_input = args.upc_input.split(',')

    # validate that UPCs are numeric
    for upc in args.upc_input:
        if not upc.isnumeric():
            parser.print_help()
            print(f'\nUPC "{upc}" is not numeric!')
            exit(1)

    # validate file type selection
    if not args.types:
        parser.print_help()
        print('\nMust select at least one asset type!')
        exit(1)

    return args


def generate_filename(row):
    """Output filename by db row."""
    out_dir = f'/tmp/output/{row["UPC"]}'
    Path(out_dir).mkdir(parents=True, exist_ok=True)
    name = '_'.join(
        str(x) for x in
        [
            row['UPC'],
            row['TRACK_VOLUME'],
            row['TRACK_INDEX'],
            row['ASSET_FINAL_ID']
        ]
        if x is not None
    )
    extension = row['S3_FILENAME'].split('.')[-1]
    return f'{out_dir}/{name}.{extension}'


async def process_queue(queue):
    """Pop from queue and do work."""
    while not queue.empty():
        data = await queue.get()
        await download_fileobj(
            data[0],
            data[1]['S3_BUCKET'],
            data[1]['S3_FILENAME'],
            generate_filename(data[1])
        )
        queue.task_done()


async def download_fileobj(s3_session, bucket, path, output_filename):
    """Download object from S3 to file."""
    async with s3_session.client('s3') as client:
        print(bucket, path, output_filename)
        with open(output_filename, 'wb+') as data:
            await client.download_fileobj(bucket, path, data)