import boto3
import argparse
import json
import datetime
import csv

from sys import exit as sysexit
from helpers.general_use import dedupe_list

def json_serial_datetime(obj):
    """
        JSON serializer for datetime objects not
        serializable by default json code

        :param obj:
        :return:
    """


    if isinstance(obj, datetime.datetime):
        serial = obj.isoformat()
        return serial
    raise TypeError ("Type not serializable")


def describeTable(dynamodb_resource):
    """
        Returns a description of the table

    :param dynamodb_resource: a dynamo db client object
    """

    table_desc = dynamodb_resource.describe_table(TableName=program.table)

    print(json.dumps(table_desc['Table'], default=json_serial_datetime))


def scanDynamoDB(dynamodb_resource, table_name, list):
    """
    This function chunks and iterates over a passed list, grabbing 100 rows at a time.

    :param dynamodb_resource:
    :param table_name:
    :param list:
    :return:
    """
    final_table = []

    chunks = [list[x:x + 100] for x in range(0, len(list), 100)]

    for chunk in chunks:
        response = dynamodb_resource.batch_get_item(
            RequestItems={
                table_name:
                    {'Keys': chunk}
            }
        )

        items = response['Responses'].get(table_name)
        #print(items)

        for item in items:
            isrc_row = pivot_array_on_isrc(item)  # Print out the subset of results.

            final_table.append(create_table(isrc_row))

    return final_table


def pivot_array_on_isrc(items):
    """
    Transforms the passed nested DynamoDB dict structure into a 2d
    array-like dict.
    This outputs the equivalent of 1 or more rows of 2D data based on the input
    structure.

    :param items: single nested structure response from a DynamoDB query
    :return:
    """

    # prep some vars
    isrc_obj = {}
    values = []
    isrc = items['isrc']['S']
    isrc_obj[isrc] = {}
    isrc_obj[isrc]['locked'] = []
    terr_list = items['territories']['M']
    locked_terr_list = items['locked_territories']['M']

    # Loop through territory list for a single ISRC
    for terr, tuid_obj in terr_list.items():
        # Reset check var
        new_tuid = True

        # Grab tuid
        tuid = tuid_obj['M']['tuid']['N']

        # Check for existing mapped tuids.
        for check_tuid in isrc_obj[isrc]:
            if new_tuid and check_tuid == tuid:
                new_tuid = False

        if new_tuid:
            isrc_obj[isrc][tuid] = []
            new_tuid = False

        isrc_obj[isrc][tuid].append(terr)

    for key, val in locked_terr_list.items():
        isrc_obj[isrc]['locked'].append(key)

        # value = JSON.stringify(items[index][header].M)

    return isrc_obj


def create_table(isrc_list_obj):
    """
    Formats the passed nested dict structure into the a pure 2d table for processing as csv

    :param isrc_list_obj:
    :return:
    """
    output_table = []
    output_row = {}

    for isrc, tuid_list in isrc_list_obj.items():
        written = False

        output_row['locked'] = ",".join(sorted(tuid_list['locked']))
        output_row['isrc'] = isrc

        # Build
        for tuid, terr_list in tuid_list.items():
            if tuid != 'locked':
                output_row['tuid'] = tuid
                output_row['territories'] = ",".join(sorted(terr_list))

                output_table.append(dict(output_row))
                written = True

        if not written:
            output_table.append(dict(output_row))

    return output_table


if __name__ == '__main__':
    headers = []
    send_expression = []

    # Setup commandline args
    parser = argparse.ArgumentParser(
        description='Dump DynamoDB table, based on a list of keys',
        add_help=False)
    parser.add_argument(
        '-t', '--table', help='Add the table you want to output to csv.',
        type=str)
    parser.add_argument(
        '-p', '--profile', help='Choose the AWS profile you want.',
        type=str)
    parser.add_argument(
        '-i', '--input_file', help='A file containing the list of ISRC\'s to'
                                   'filter the dump.', type=str)
    parser.add_argument(
        '-o', '--output_file', help='The output file to which to write.', type=str)
    parser.add_argument(
        '-d', '--describe', help='Only describe the table, and quit.',
        action='store_true', default=False)

    # Parse commandline args
    program = parser.parse_args()


    if not program.table:
        print("You must specify a table");
        sysexit(1)

    if program.profile:
        boto3.setup_default_session(profile_name=program.profile)

    # Open client connection
    dynamoDB = boto3.client('dynamodb')

    # User specified a description only
    if (program.describe):
        describeTable(dynamoDB)
    else:  # Full query requested

        if program.input_file:
            with open(program.input_file) as f:
                # Create clean list of DEDUPED isrc's
                filter_isrc_list = dedupe_list([line.rstrip() for line in f])

            # Convert to Dynamo request structure
            send_expression = [{'isrc': {'S': val}} for val in
                               filter_isrc_list]
        else:
            print("You must specify an input filename.")
            sysexit(1)

        # HERE'S THE MEAT
        final_table = scanDynamoDB(dynamoDB, program.table, send_expression)

        if not program.output_file:
            program.output_file = 'dynamo_dump.csv'
            print('No output file specified. Using \'dynamo_dump.csv\'')

        # Output table to csv
        with open(program.output_file, 'w', encoding='utf8') as f:
            writer = csv.DictWriter(f, final_table[0][0].keys())
            writer.writeheader()
            for slice in final_table:
                for row in slice:
                    writer.writerow(row)

        print ('Output written to ', program.output_file)