"""Tasks to invoke the various Lambdas composing the `adjustments_file_generate` DAG."""

from hooks.lambda_hook import OrchLambdaHook
from lib import config
from tasks.adjustment_file_generate import helpers
from tasks.adjustment_file_generate.states import success_generate_task


def invoke_generate_task(dag_run, **kwargs):
    """Invoke the `generate_flowthrough_adjustments` Lambda.

    NOTE: This task is being used as a `ShortCircuitOperator`, it needs
    to return `True` to execute the downstream tasks, or `False` to skip them.
    """
    event = helpers.get_event_from_params(dag_run, **kwargs)

    # NOTE: We're specifying a different `aws_conn_id` here because the QA
    # variant of the Lambda is hosted in the dedicated Accounting AWS account.
    if config.OWS_ENV.lower() == 'qa':
        hook = OrchLambdaHook(
            config.GENERATE_FLOWTHROUGH_ADJUSTMENTS_LAMBDA_NAME,
            aws_conn_id='accounting_qa'
        )
    else:
        hook = OrchLambdaHook(config.GENERATE_FLOWTHROUGH_ADJUSTMENTS_LAMBDA_NAME)

    response = hook.invoke_lambda(event.to_json())

    if not response.function_response.succeeded:
        raise Exception(response.error_message)

    status_ok = response.function_response.payload.get('status') == 'OK'

    # NOTE: The Lambda might return a `NO_RECORDS` status, meaning that there
    # were no adjustments to generate. In this case we want to skip the
    # remaining tasks but not fail the DAG.
    if not status_ok:
        print('No adjustments generated. Marking the task as successfull and stopping the DAG.')  # noqa: E501
        success_generate_task(dag_run, **kwargs)
        return False

    return True


def invoke_validate_task(dag_run, **kwargs):
    """Invoke the `adjustments_json_validation` Lambda."""
    event = helpers.get_event_from_params(dag_run, **kwargs)

    # NOTE: We're specifying a different `aws_conn_id` here because the QA
    # variant of the Lambda is hosted in the dedicated Accounting AWS account.
    if config.OWS_ENV.lower() == 'qa':
        hook = OrchLambdaHook(
            config.ADJUSTMENTS_JSON_VALIDATION_LAMBDA_NAME,
            aws_conn_id='accounting_qa'
        )
    else:
        hook = OrchLambdaHook(config.ADJUSTMENTS_JSON_VALIDATION_LAMBDA_NAME)

    response = hook.invoke_lambda(event.to_json())

    if not response.function_response.succeeded:
        raise Exception(response.error_message)

    status_ok = response.function_response.payload.get('status') == 'OK'

    # NOTE: If the Lambda returns some validation errors,
    # we want to fail the DAG.
    if not status_ok:
        raise Exception('Validation Errors')


def invoke_import_task(dag_run, **kwargs):
    """Invoke the `adjustments_json_import` Lambda."""
    event = helpers.get_event_from_params(dag_run, **kwargs)

    # NOTE: We're specifying a different `aws_conn_id` here because the QA
    # variant of the Lambda is hosted in the dedicated Accounting AWS account.
    if config.OWS_ENV.lower() == 'qa':
        hook = OrchLambdaHook(
            config.ADJUSTMENTS_JSON_IMPORT_LAMBDA_NAME,
            aws_conn_id='accounting_qa'
        )
    else:
        hook = OrchLambdaHook(config.ADJUSTMENTS_JSON_IMPORT_LAMBDA_NAME)

    response = hook.invoke_lambda(event.to_json())

    if not response.function_response.succeeded:
        raise Exception(response.error_message)