"""Distribution Fee ETL Tasks."""
from datetime import date

from garcon import task
from garcon_contrib.dynamo_feed_status import garcon_feed_status

from flows import art_relations
from flows import datastore
from flows import util
from flows.distribution_fee import config
from flows.distribution_fee import log
from flows.distribution_fee import queries
from flows.distribution_fee import status
from flows.distribution_fee import util as etl_util


@task.decorate(timeout=120)
def create_temp_table(activity, correlation_id, create_statement, table_name):
    """Create temp tables for the ETL.

    Args:
        activity (ActivityWorker): activity worker.
        correlation_id (str): etl correlation ID.
        create_statement (str): CREATE TABLE query to run.
        table_name (str): temp table name format string.

    Returns:
        dict: table name.
    """
    correlation_hex = util.correlation_id_hex(correlation_id)
    table_name = table_name.format(cid=correlation_hex)
    sql = create_statement.format(table_name=table_name)
    datastore.execute(sql)

    log.update_status(correlation_id, status.TEMP_TABLES_CREATED)

    return {'name': table_name}


@task.decorate(timeout=3600)
def load_vendor_contract(
        activity, call_query, correlation_id, insert_query, table_name, upcs):
    """Load vendor contracts for all upcs for calculation.

    Args:
        activity (ActivityWorker): activity worker.
        call_query (str): stored procedure to get vendor contract id.
        correlation_id (str): etl correlation ID.
        insert_query (str): query to populate temp table with contract data.
        table_name (str): temp table name format string.
        upcs (list(str)): list of upcs to process.
    """
    sql = insert_query.format(table_name=table_name)
    insert_rows = []
    vendor_contracts = {}
    upc_vendor_ids_iter = util.get_upc_vendor_id(upcs)
    for upc_vendor_ids in upc_vendor_ids_iter:
        for upc, vendor_id in upc_vendor_ids.items():
            if vendor_id not in vendor_contracts:
                sp_call = art_relations.query(call_query, (vendor_id,))
                call_result = sp_call.fetchone()
                if call_result is None:
                    continue
                vendor_contracts[vendor_id] = call_result[0]
            insert_rows.append((vendor_contracts[vendor_id], vendor_id, upc))
    datastore.executemany(sql, insert_rows)

    log.update_status(correlation_id, status.LOADED_CONTRACT_IDS)


@task.decorate(timeout=1800)
def load_dist_fee_regular(activity, contract_table, correlation_id, fee_table):
    """Load vendor contracts for all upcs for calculation.

    Args:
        activity (ActivityWorker): activity worker.
        contract_table (str): temp vendor contract table to populate from.
        correlation_id (str): etl correlation ID.
        fee_table (str): temp fee table to populate.
    """
    vendor_contract_ids, vendor_upc_map = etl_util.get_etl_vendor_contracts(
        contract_table)

    fee_rows = []
    contract_queries = queries.get_select_contract_fee_sql(
        vendor_contract_ids, 'regular')
    for contract_query in contract_queries:
        dist_fees_results = art_relations.query(contract_query)
        dist_fees = dist_fees_results.fetchall()
        for vendor_id, split in dist_fees:
            for upc in vendor_upc_map[vendor_id]:
                fee_rows.append([split, upc])

    insert_query = queries.INSERT_FEE_REGULAR.format(table_name=fee_table)
    datastore.executemany(insert_query, fee_rows)
    log.update_status(correlation_id, status.DIST_FEE_REGULAR_LOADED)


@task.decorate(timeout=1800)
def load_dist_fee_territory(
        activity, contract_table, correlation_id, fee_table):
    """Load vendor territory contracts for all upcs for calculation.

    Args:
        activity (ActivityWorker): activity worker.
        contract_table (str): temp vendor contract table to populate from.
        correlation_id (str): etl correlation ID.
        fee_table (str): temp fee table to populate.
    """
    vendor_contract_ids, vendor_upc_map = etl_util.get_etl_vendor_contracts(
        contract_table)

    fee_rows = []
    contract_queries = queries.get_select_contract_fee_sql(
        vendor_contract_ids, 'territory')
    for contract_query in contract_queries:
        dist_fees_results = art_relations.query(contract_query)
        dist_fees = dist_fees_results.fetchall()
        for vendor_id, country_id, split in dist_fees:
            for upc in vendor_upc_map[vendor_id]:
                fee_rows.append([country_id, split, upc])

    insert_query = queries.INSERT_FEE_TERRITORY.format(table_name=fee_table)
    datastore.executemany(insert_query, fee_rows)
    log.update_status(correlation_id, status.DIST_FEE_TERRITORY_LOADED)


@task.decorate(timeout=7200)
def calculate_client_amount(
        activity, correlation_id, fee_dms_table, fee_ter_table, fee_reg_table,
        upcs):
    """Calculate and store client amounts.

    In a single transaction, this removes all client amount rows for the
    specified UPCs regardless of transaction or revenue type, and stores new
    calculations for them.

    Args:
        activity (ActivityWorker): activity worker.
        correlation_id (str): etl correlation ID.
        fee_dms_table (str): temp table with dms fee data.
        fee_ter_table (str): temp table with territory fee data.
        fee_reg_table (str): temp table with regular fee data.
        upcs (list(str)): list of upcs to process.
    """
    with datastore.context() as (cursor, _):  # noqa
        delete_sqls = queries.get_delete_client_amount_sql(upcs)
        for delete_sql in delete_sqls:
            cursor.execute(delete_sql)
        for table, column in config.REVENUE_TABLE_COLUMNS.items():
            insert_sqls = queries.get_insert_client_amount_sql(
                table, column, fee_dms_table, fee_ter_table, fee_reg_table,
                upcs)
            for insert_sql in insert_sqls:
                cursor.execute(insert_sql)

    log.update_status(correlation_id, status.CLIENT_AMOUNT_INSERTED)


@task.decorate(timeout=900)
def update_distribution_fee_table(
        activity, correlation_id, fee_dms_table, fee_ter_table, fee_reg_table,
        upcs):
    """Update distribution fee table.

    Args:
        activity (ActivityWorker): activity worker.
        correlation_id (str): etl correlation ID.
        fee_dms_table (str): temp table with dms fee data.
        fee_ter_table (str): temp table with territory fee data.
        fee_reg_table (str): temp table with regular fee data.
        upcs (list(str)): list of upcs to process.
    """
    with datastore.context() as (cursor, _):
        delete_sqls = queries.get_delete_distribution_table_data_sql(upcs)
        for delete_sql in delete_sqls:
            cursor.execute(delete_sql)

        cursor.execute(
            queries.get_insert_distribution_table_data_sql(
                fee_dms_table, 'dms'))
        cursor.execute(
            queries.get_insert_distribution_table_data_sql(
                fee_ter_table, 'territory'))
        cursor.execute(
            queries.get_insert_distribution_table_data_sql(
                fee_reg_table, 'regular'))

        log.update_status(
            correlation_id, status.DISTRIBUTION_FEE_TABLE_UPDATED)


@task.decorate(timeout=120)
def drop_temp_table(activity, correlation_id, drop_statement, table_name):
    """Drop the ETL temp tables.

    Args:
        activity (ActivityWorker): activity worker.
        correlation_id (str): etl correlation ID.
        drop_statement (str): DROP TABLE query format string.
        table_name (str): table name to drop.
    """
    sql = drop_statement.format(table_name=table_name)
    datastore.execute(sql)

    log.update_status(correlation_id, status.TEMP_TABLES_DROPPED)


@task.decorate(timeout=120)
def set_dynamo_status(activity, correlation_id):
    """Set DynamoDB statuses.

    Args:
        activity (Activity): SWF Activity.
        correlation_id (str): etl correlation ID.
    """
    report_date = date.today().strftime('%Y-%m-%d')
    garcon_feed_status.set_overall_status(
        config.SWF_WORKFLOW_NAME, report_date,
        garcon_feed_status.STATUS_INGESTED)

    log.update_status(correlation_id, status.COMPLETED)


@task.decorate(timeout=120)
def send_sns_notification(activity, correlation_id, upcs):
    """Send all listeners a message that the ETL has completed.

    Args:
        activity (ActivityWorker): activity worker.
        correlation_id (str): etl correlation ID.
        upcs (list): upcs for unload query.
    """
    report_date = date.today().strftime('%Y-%m-%d')
    sns_correlation_id = correlation_id + '.1'
    etl_util.send_sns_notification(sns_correlation_id, report_date, upcs)