"""Tasks to update the Abacus state of the various actions composing the `adjustment_file_generate` DAG."""  # noqa: E501

from lib import config
from lib.constants import (
    ABACUS_STATE_STATUSES,
    AUTO_GENERATED_STATEMENT_PERIOD_ADJUSTMENT_FILE_ACTIONS,
)
from lib.utils import ows
from tasks.adjustment_file_generate import helpers


# NOTE: We're sending the DAG name as the Orchard identity header in order to
# be able to toggle the `abacus_flowthrough_automation` feature flag, which
# dictates the list of actions available for
# the `statement_period_adjustment_file` target type.
DEFAULT_UPDATE_HEADERS = {
    'Orchard-Identity-Id': config.ADJUSTMENT_FILE_GENERATE_DAG_NAME
}


def start_generate_task(dag_run, **kwargs):
    """Mark the "generate" action as running."""
    event = helpers.get_event_from_params(dag_run, **kwargs)

    action_name = (
        AUTO_GENERATED_STATEMENT_PERIOD_ADJUSTMENT_FILE_ACTIONS.UPLOAD_FILE
    )

    abacus_state = helpers.get_abacus_state(action_name, event.target_id)

    ows.update_abacus_state(
        abacus_state.get('abacus_state_id'),
        body={
            'action_status': ABACUS_STATE_STATUSES.RUNNING
        },
        headers=DEFAULT_UPDATE_HEADERS
    )


def success_generate_task(dag_run, **kwargs):
    """Mark the "generate" action as complete."""
    event = helpers.get_event_from_params(dag_run, **kwargs)

    action_name = (
        AUTO_GENERATED_STATEMENT_PERIOD_ADJUSTMENT_FILE_ACTIONS.UPLOAD_FILE
    )

    abacus_state = helpers.get_abacus_state(action_name, event.target_id)

    ows.update_abacus_state(
        abacus_state.get('abacus_state_id'),
        body={
            'action_status': ABACUS_STATE_STATUSES.COMPLETE
        },
        headers=DEFAULT_UPDATE_HEADERS
    )


def failure_generate_task(dag_run, **kwargs):
    """Mark the "generate" action as failed.

    NOTE: This task is being used as a `ShortCircuitOperator`, it always
    returns `False` to skip the downstream tasks.
    """
    event = helpers.get_event_from_params(dag_run, **kwargs)

    action_name = (
        AUTO_GENERATED_STATEMENT_PERIOD_ADJUSTMENT_FILE_ACTIONS.UPLOAD_FILE
    )

    abacus_state = helpers.get_abacus_state(action_name, event.target_id)

    ows.update_abacus_state(
        abacus_state.get('abacus_state_id'),
        body={
            'action_status': ABACUS_STATE_STATUSES.ERROR
        },
        headers=DEFAULT_UPDATE_HEADERS
    )

    return False


def start_validate_task(dag_run, **kwargs):
    """Mark the "validate" action as running."""
    event = helpers.get_event_from_params(dag_run, **kwargs)

    action_name = (
        AUTO_GENERATED_STATEMENT_PERIOD_ADJUSTMENT_FILE_ACTIONS.VALIDATE_FILE
    )

    abacus_state = helpers.get_abacus_state(action_name, event.target_id)

    ows.update_abacus_state(
        abacus_state.get('abacus_state_id'),
        body={
            'action_status': ABACUS_STATE_STATUSES.RUNNING
        },
        headers=DEFAULT_UPDATE_HEADERS
    )


def success_validate_task(dag_run, **kwargs):
    """Mark the "validate" action as complete."""
    event = helpers.get_event_from_params(dag_run, **kwargs)

    action_name = (
        AUTO_GENERATED_STATEMENT_PERIOD_ADJUSTMENT_FILE_ACTIONS.VALIDATE_FILE
    )

    abacus_state = helpers.get_abacus_state(action_name, event.target_id)

    ows.update_abacus_state(
        abacus_state.get('abacus_state_id'),
        body={
            'action_status': ABACUS_STATE_STATUSES.COMPLETE
        },
        headers=DEFAULT_UPDATE_HEADERS
    )


def failure_validate_task(dag_run, **kwargs):
    """Mark the "validate" action as failed.

    NOTE: This task is being used as a `ShortCircuitOperator`, it always
    returns `False` to skip the downstream tasks.
    """
    event = helpers.get_event_from_params(dag_run, **kwargs)

    action_name = (
        AUTO_GENERATED_STATEMENT_PERIOD_ADJUSTMENT_FILE_ACTIONS.VALIDATE_FILE
    )

    abacus_state = helpers.get_abacus_state(action_name, event.target_id)

    ows.update_abacus_state(
        abacus_state.get('abacus_state_id'),
        body={
            'action_status': ABACUS_STATE_STATUSES.ERROR
        },
        headers=DEFAULT_UPDATE_HEADERS
    )

    return False


def start_import_task(dag_run, **kwargs):
    """Mark the "import" action as running."""
    event = helpers.get_event_from_params(dag_run, **kwargs)

    action_name = (
        AUTO_GENERATED_STATEMENT_PERIOD_ADJUSTMENT_FILE_ACTIONS.IMPORT_FILE
    )

    abacus_state = helpers.get_abacus_state(action_name, event.target_id)

    ows.update_abacus_state(
        abacus_state.get('abacus_state_id'),
        body={
            'action_status': ABACUS_STATE_STATUSES.RUNNING
        },
        headers=DEFAULT_UPDATE_HEADERS
    )


def success_import_task(dag_run, **kwargs):
    """Mark the "import" action as complete."""
    event = helpers.get_event_from_params(dag_run, **kwargs)

    action_name = (
        AUTO_GENERATED_STATEMENT_PERIOD_ADJUSTMENT_FILE_ACTIONS.IMPORT_FILE
    )

    abacus_state = helpers.get_abacus_state(action_name, event.target_id)

    ows.update_abacus_state(
        abacus_state.get('abacus_state_id'),
        body={
            'action_status': ABACUS_STATE_STATUSES.COMPLETE
        },
        headers=DEFAULT_UPDATE_HEADERS
    )


def failure_import_task(dag_run, **kwargs):
    """Mark the "import" action as failed.

    NOTE: This task is being used as a `ShortCircuitOperator`, it always
    returns `False` to skip the downstream tasks.
    """
    event = helpers.get_event_from_params(dag_run, **kwargs)

    action_name = (
        AUTO_GENERATED_STATEMENT_PERIOD_ADJUSTMENT_FILE_ACTIONS.IMPORT_FILE
    )

    abacus_state = helpers.get_abacus_state(action_name, event.target_id)

    ows.update_abacus_state(
        abacus_state.get('abacus_state_id'),
        body={
            'action_status': ABACUS_STATE_STATUSES.ERROR
        },
        headers=DEFAULT_UPDATE_HEADERS
    )

    return False