"""Airflow connector."""

import base64
import json
import logging
import uuid

import boto3
import requests
from werkzeug.exceptions import BadRequest

from abacus_event.config import Config
from abacus_event.constants.constants import DAG_IDS, EVENT_NAMES, PARALLEL_DAG_RUNS
from abacus_event.models.abacus_event import AbacusEvent
from abacus_event.schemas.abacus_event import AbacusEventSchema

abacus_event_schema = AbacusEventSchema()


def handle_event_actions(abacus_event):
    """Trigger a corresponding DAG, if one exists."""
    event_handler = event_handlers.get(abacus_event.event_name)
    if not event_handler:
        return

    dag_id, conf = event_handler(abacus_event)
    ensure_dag_not_running_in_aws(dag_id)
    trigger_dag_from_aws(dag_id, conf)


def ensure_dag_not_running_in_aws(dag_id):
    """Raise if specified DAG is currently in progress in AWS Airflow."""
    request_body = f'dags list-runs --state running -o json --dag-id {dag_id}'

    try:
        result = _aws_airflow_request(request_body)
    except BadRequest:
        # Most likely error related to the issue above
        return True
    else:
        dag_runs = json.loads(result)
        if not dag_runs:
            return True
        elif dag_id in PARALLEL_DAG_RUNS and len(dag_runs) < PARALLEL_DAG_RUNS[dag_id]:
            # This dag is allowed to have multiple concurrent runs
            return True
        logging.error('DAG %s already in progress: %s', dag_id, dag_runs)
        raise BadRequest(f"DAG '{dag_id}' already in progress.")


def trigger_dag_from_aws(dag_id, config=''):
    """Trigger a DAG from Managed Airflow Environment in AWS."""
    dag_config = json.dumps(config)
    dag_run_id = ''
    if dag_id in PARALLEL_DAG_RUNS:
        # Assign a unique dag run id in the event of colliding timestamp
        dag_run_id = f'-r {dag_id}_{uuid.uuid4()} '

    body = f"""dags trigger {dag_run_id}-c '{dag_config}' {dag_id}"""
    return _aws_airflow_request(body)


def _aws_airflow_request(request_body):
    """Request Managed Airflow Environment in AWS."""
    client = boto3.client('mwaa')
    token = client.create_cli_token(Name=Config.AWS_AIRFLOW_ENVIRONMENT_NAME)
    url = 'https://{0}/aws_mwaa/cli'.format(token['WebServerHostname'])
    headers = {
        'Authorization': 'Bearer ' + token['CliToken'],
        'Content-Type': 'text/plain',
    }
    result = requests.post(url, data=request_body, headers=headers)
    if result.status_code not in range(200, 300):
        raise BadRequest(f'AWS Airflow API Error: ({result.status_code}) {result.text}')

    error_message = result.json()['stderr']
    output_message = result.json()['stdout']
    if error_message:
        sanitized_message = _sanitize_error_message(error_message)
        if sanitized_message:
            raise BadRequest(sanitized_message)
    return base64.b64decode(output_message).decode('utf-8')


def _sanitize_error_message(error_message):
    """Remove deprecation warnings from output."""
    decoded_messages = base64.b64decode(error_message).decode('utf-8').split('\n')
    messages = list()

    for message in decoded_messages:
        if not message:
            continue
        if 'warning' in message.lower() or 'getattr(module, class_name)' in message:
            continue
        messages.append(message)

    return '\n'.join(messages)


def _accounting_period_close_dag_config(abacus_event: AbacusEvent) -> tuple:
    """Return the dag name and config for the accounting period close workflow."""
    return DAG_IDS.ACCOUNTING_PERIOD_CLOSE, abacus_event_schema.dump(abacus_event)


def _accounting_period_mechanicals_dag_config(abacus_event):
    """Return the dag name and config for accounting period mechanicals workflow."""
    return DAG_IDS.ACCOUNTING_PERIOD_MECHANICALS, abacus_event_schema.dump(abacus_event)


def _accounting_period_sales_approve_dag_config(abacus_event):
    """Return the dag name and config for accounting period sales approve workflow."""
    return DAG_IDS.ACCOUNTING_PERIOD_SALES_APPROVE, abacus_event_schema.dump(
        abacus_event
    )


def _accounting_run_calculate_dag_config(abacus_event):
    """Return the dag name and config for the accounting run calculate event."""
    return DAG_IDS.ACCOUNTING_RUN_CALCULATE, abacus_event_schema.dump(abacus_event)


def _accounting_run_calculate_nr_dag_config(abacus_event):
    """Return the dag name and config for the accounting run calculate NR event."""
    return DAG_IDS.ACCOUNTING_RUN_CALCULATE_NR, abacus_event_schema.dump(abacus_event)


def _accounting_run_commit_dag_config(abacus_event):
    """Return the dag name and config for the accounting run commit event."""
    return DAG_IDS.ACCOUNTING_RUN_COMMIT, abacus_event_schema.dump(abacus_event)


def _adjustment_file_generate_dag_config(abacus_event):
    """Return the dag name and config for adjustment file generate workflow."""
    return DAG_IDS.ADJUSTMENT_FILE_GENERATE, abacus_event_schema.dump(abacus_event)


def _adjustment_file_import_dag_config(abacus_event):
    """Return the dag name and config for adjustment file upload workflow."""
    return DAG_IDS.ADJUSTMENT_FILE_IMPORT, abacus_event_schema.dump(abacus_event)


def _adjustment_file_upload_dag_config(abacus_event):
    """Return the dag name and config for adjustment file upload workflow."""
    return DAG_IDS.ADJUSTMENT_FILE_UPLOAD, abacus_event_schema.dump(abacus_event)


def _apply_adjustments_dag_config(abacus_event):
    """Return the dag name and config for adjustments apply workflow."""
    return DAG_IDS.APPLY_PENDING_ADJUSTMENTS, abacus_event_schema.dump(abacus_event)


def _payments_generate_dag_config(abacus_event):
    """Return the dag name and config for the create payment workflow."""
    return DAG_IDS.PAYMENTS_GENERATE, abacus_event_schema.dump(abacus_event)


def _payments_generate_export_dag_config(abacus_event):
    """Return the dag name and config for the generate payment export workflow."""
    return DAG_IDS.PAYMENTS_GENERATE_EXPORT, abacus_event_schema.dump(abacus_event)


def _payments_upload_approval_dag_config(abacus_event):
    """Return the dag name and config for the upload payment approval workflow."""
    return DAG_IDS.PAYMENTS_UPLOAD_APPROVAL, abacus_event_schema.dump(abacus_event)


def _payoneer_payments_payout_dag_config(abacus_event):
    """Return the dag name and config for payoneer payments workflow."""
    return DAG_IDS.PAYONEER_PAYMENTS_PAYOUT, abacus_event_schema.dump(abacus_event)


def _reserves_release_dag_config(abacus_event):
    """Return the dag name and config for reserves release workflow."""
    return DAG_IDS.RESERVES_RELEASE, abacus_event_schema.dump(abacus_event)


def _sales_approve_dag_config(abacus_event):
    """Return the dag name and config for the sales get eligible event."""
    return DAG_IDS.SALES_APPROVE, abacus_event_schema.dump(abacus_event)


def _sales_get_eligible_dag_config(abacus_event):
    """Return the dag name and config for the sales get eligible event."""
    return DAG_IDS.SALES_GET_ELIGIBLE, abacus_event_schema.dump(abacus_event)


def _sales_ingest_dag_config(abacus_event):
    """Return the dag name and config for the sales ingest event."""
    return DAG_IDS.SALES_INGEST, abacus_event_schema.dump(abacus_event)


# Mapping of event names to handlers. Handlers should return:
#     Tuple(String dag_name, Dict config)
event_handlers = {
    EVENT_NAMES.ACCOUNTING_PERIOD_CLOSE: _accounting_period_close_dag_config,
    EVENT_NAMES.ACCOUNTING_PERIOD_APPROVE_SALES_FILES: _accounting_period_sales_approve_dag_config,
    EVENT_NAMES.ACCOUNTING_PERIOD_MECHANICALS: _accounting_period_mechanicals_dag_config,
    EVENT_NAMES.ACCOUNTING_RUN_CALCULATE: _accounting_run_calculate_dag_config,
    EVENT_NAMES.ACCOUNTING_RUN_CALCULATE_NR: _accounting_run_calculate_nr_dag_config,
    EVENT_NAMES.ACCOUNTING_RUN_COMMIT: _accounting_run_commit_dag_config,
    EVENT_NAMES.ADJUSTMENT_FILE_UPLOAD: _adjustment_file_upload_dag_config,
    EVENT_NAMES.ADJUSTMENT_FILE_WORKSHEET_IMPORT: _adjustment_file_import_dag_config,
    EVENT_NAMES.APPLY_PENDING_ADJUSTMENTS: _apply_adjustments_dag_config,
    EVENT_NAMES.GENERATE_EXPORT: _payments_generate_export_dag_config,
    EVENT_NAMES.GENERATE_FLOWTHROUGH_ADJUSTMENTS: _adjustment_file_generate_dag_config,
    EVENT_NAMES.PAYMENTS_GENERATE: _payments_generate_dag_config,
    EVENT_NAMES.RELEASE_RESERVES: _reserves_release_dag_config,
    EVENT_NAMES.SALES_APPROVE: _sales_approve_dag_config,
    EVENT_NAMES.SALES_GET_ELIGIBLE: _sales_get_eligible_dag_config,
    EVENT_NAMES.SALES_INGEST_DISTRO: _sales_ingest_dag_config,
    EVENT_NAMES.SALES_INGEST_NR: _sales_ingest_dag_config,
    EVENT_NAMES.SEND_PAYMENTS: _payoneer_payments_payout_dag_config,
    EVENT_NAMES.UPLOAD_APPROVAL: _payments_upload_approval_dag_config,
}