"""Lambda to trigger distribution_fee ETL when other ETLs have completed."""
from datetime import datetime
import json
from os import environ

import boto3

import mysql


SWF_DOMAIN = environ.get('SWF_DOMAIN') or 'dev'
SWF_REGION_NAME = 'us-east-1'
SWF_WORKFLOW_ID_PATTERN = 'ows_ft_etl_distribution_fee_{cid}'
SWF_WORKFLOW_NAME = 'ows_ft_etl_distribution_fee'
SWF_WORKFLOW_TIMEOUT = '7200'
SWF_WORKFLOW_VERSION = '1.0'
VALID_PRECEDING_SOURCES = {
    'cable_calculation', 'digital', 'theatrical', 'theatrical_cuts'}


def lambda_handler(event, context):
    """Lambda starting point."""
    print('Event => ' + event['Records'][0]['Sns']['Message'])
    try:
        message = json.loads(event['Records'][0]['Sns']['Message'])
    except ValueError:
        print('Invalid Json message')
        return

    # check if the caller is valid ETL.
    if not message.get('correlation_id') \
            or message.get('source') not in VALID_PRECEDING_SOURCES:
        return

    correlation_id = message.get('correlation_id')
    upcs = message.get('upcs')
    swf_context = {'correlation_id': correlation_id, 'upcs': upcs}

    print('Triggering dist_fee etl with => ' + json.dumps(swf_context))
    if not create_log(**swf_context):
        print('Failed to create log. Terminating.')
        return

    create_etl_context(correlation_id, upcs)
    swf_context['upcs'] = 'upcs'
    swf = boto3.client('swf', region_name=SWF_REGION_NAME)
    response = swf.start_workflow_execution(
        input=json.dumps(swf_context),
        domain=SWF_DOMAIN,
        taskList={'name': SWF_WORKFLOW_NAME},
        executionStartToCloseTimeout=SWF_WORKFLOW_TIMEOUT,
        workflowId=SWF_WORKFLOW_ID_PATTERN.format(cid=correlation_id),
        workflowType={
            'name': SWF_WORKFLOW_NAME,
            'version': SWF_WORKFLOW_VERSION})

    add_run_id_to_log(correlation_id, response.get('runId'))
    return response


def create_log(correlation_id, upcs):
    """Create a new ETL log entry.

    Args:
        correlation_id (str): Correlation ID for ETL process.
        upcs (tuple): upc strings.
    """
    sql = """
    INSERT
        distribution_fee_etl_log
        (correlation_id, date_end, date_start, etl_start, etl_status, upcs)
    VALUES
        (
            %(correlation_id)s, null, null, %(etl_start)s,
            %(etl_status)s, %(upcs)s)"""

    params = {
        'correlation_id': correlation_id.split('.')[0],
        'etl_start': datetime.now(),
        'etl_status': 'STARTED', 'upcs': json.dumps(upcs)}
    try:
        print('Creating DB log for ' + correlation_id)
        return mysql.execute(sql, params)
    except Exception as e:
        print(e)
        return False


def add_run_id_to_log(correlation_id, workflow_run_id):
    """Update ETL log with run ID from AWS SWF.

    Args:
        correlation_id (str): Correlation ID for ETL process.
        workflow_run_id (str): SWF execution run id.
    """
    sql = """
    UPDATE
        distribution_fee_etl_log
    SET
        workflow_run_id = %(workflow_run_id)s
    WHERE
        correlation_id = %(correlation_id)s"""

    params = {
        'correlation_id': correlation_id.split('.')[0],
        'workflow_run_id': workflow_run_id}

    try:
        print('Updating DB log for ' + correlation_id)
        return mysql.execute(sql, params)
    except Exception as e:
        print(e)
        return False


def create_etl_context(correlation_id, upcs):
    """Create DatabaseParam style upcs context if it doesn't exist.

    This is backwards compatible. If the upcs parameter is not a list but a
    context key, it will assume the context already exists and skip the insert.

    Args:
        correlation_id (str): correlation ID sent in lambda.
        upcs (list): list of UPCs to store.
    """
    if upcs == 'upcs':
        return

    params = {
        'context_key': 'upcs',
        'correlation_id': correlation_id.split('.')[0],
        'data': json.dumps(upcs)}

    sql = """
        INSERT INTO
            etl_context (correlation_id, context_key, data)
        VALUES
            (%(correlation_id)s, %(context_key)s, %(data)s)
        ON DUPLICATE KEY UPDATE
            data = %(data)s;"""

    mysql.execute(sql, params)