"""Requests to aws airflow."""

import base64
import os

import boto3
import requests
from owsresponse import response

from core.config import Config


def _sanitize_error_message(error_message):
    """Get errors 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 run_airflow_cli_command(command: str) -> response.Response:
    """Run an arbitrary airflow CLI command."""
    airflow_env_name = (
        'prod-abacus-airflow'
        if Config.ENVIRONMENT == Config.PROD_ENVIRONMENT
        else 'qa-abacus-airflow'
    )
    client = boto3.client('mwaa')
    token = client.create_cli_token(Name=airflow_env_name)
    url = 'https://{0}/aws_mwaa/cli'.format(token['WebServerHostname'])
    headers = {
        'Authorization': 'Bearer ' + token['CliToken'],
        'Content-Type': 'text/plain',
    }
    result = requests.post(url, data=command, headers=headers)
    error_message = result.json()['stderr']
    output_message = result.json()['stdout']
    errors = None
    if error_message:
        sanitized_message = _sanitize_error_message(error_message)
        if sanitized_message:
            errors = sanitized_message
    decoded_output_message = base64.b64decode(output_message).decode('utf-8')
    return {'output': decoded_output_message, 'errors': errors}