import asyncio
import concurrent.futures
import csv

import aioboto3
import smart_open

from ssavva.sync_to_dynamodb import config
from ssavva.sync_to_dynamodb import utils


def get_products(s3_path):
    items = []
    with smart_open.smart_open(s3_path, 'r') as f:
        reader = csv.reader(f, quotechar='"', delimiter=',')
        for r in reader:
            items.append(r)
    return items

async def producer(loop, queue, bucket, path):
    """
        Producer runs synchronous get_products function in a separate
        thread by async loop in order do not block async operations.
        Results of get_products are put in async queue.
    """
    s3_keys = utils.keys_in_s3_location(bucket, path)
    for chunk in utils.chunk_list(s3_keys, 2):
        with concurrent.futures.ThreadPoolExecutor() as pool:
            futures = [
                loop.run_in_executor(pool, get_products, s3_path)
                for s3_path in chunk]
            results = await asyncio.gather(*futures)
            for result in results:
                await queue.put(result)

    await queue.put(None)


async def consumer(queue, semaphore):
    """
        Consumer waits for the new portions of data from async queue in
        infinite loop and runs async tasks. Number of simultaneous task is
        controlled by async semaphore (processing of all data simultaneously
        consumes too much memory).
    """
    while True:
        items = await queue.get()
        if items is None:
            break
        async with semaphore:
            task = asyncio.create_task(put_in_dynamo(items, semaphore))
            asyncio.ensure_future(task)


async def put_in_dynamo(items, semaphore):
    """
        Async work with DynamoDB.
    """
    await semaphore.acquire()
    try:
        async with aioboto3.resource('dynamodb') as dynamo_resource:
            table = dynamo_resource.Table(config.dynamo_table_name)
            async with table.batch_writer() as batch:
                for item in items:
                    await batch.put_item(dict(zip(config.product_fields, item)))
    except Exception as e:
        print(str(e))
    finally:
        semaphore.release()


def main(data_type):
    loop = asyncio.get_event_loop()
    # We keep buffer (queue size) as small as possible in order to save memory
    queue = asyncio.Queue(maxsize=1)

    semaphore = asyncio.BoundedSemaphore(20)

    bucket = config.s3_bucket
    path = 'dynamo_sync/{}_full/'.format(data_type)

    producer_coroutine = producer(loop, queue, bucket, path)
    consumer_coroutine = consumer(queue, semaphore)

    loop.run_until_complete(
        asyncio.gather(producer_coroutine, consumer_coroutine))


if __name__ == '__main__':
    main('products')