"""Lambda manage-kafka-connector module."""
import json

import boto3
import requests
from kafka.admin import ConfigResource
from kafka.admin import KafkaAdminClient
from lambdacommon.common_config import logger
import sentry_sdk
from sentry_sdk.integrations.aws_lambda import AwsLambdaIntegration

import config

if config.SENTRY_DSN:
    sentry_sdk.init(
        dsn=config.SENTRY_DSN,
        environment=config.ENVIRONMENT,
        integrations=[AwsLambdaIntegration(timeout_warning=True)]
    )


def assume_target_account_role(account_id, role):
    """
    Assume IAM role in target account.

    Args:
    account_id (str): AWS account ID

    Returns:
    dict: dict of AWS credentials
    """
    client = boto3.client('sts')
    assume_role_response = client.assume_role(
        RoleArn=f'arn:aws:iam::{account_id}:role/{role}',
        RoleSessionName=config.SERVICE_NAME,
        ExternalId=config.EXTERNAL_ID,
        DurationSeconds=3600
    )
    credentials = assume_role_response['Credentials']

    return credentials


def describe_service(ecs_client, cluster_name, service_name):
    """Return service details."""
    response = ecs_client.describe_services(
        cluster=cluster_name,
        services=[service_name]
    )

    services = response['services']
    if not services:
        raise Exception(
            f'Service {service_name} does not exist in cluster {cluster_name}')

    service = services[0]
    return service


def scale_service(
        ecs_client, cluster_name, service_name, desired_count, wait=True):
    """Scale an ECS service in or out."""
    service = describe_service(ecs_client, cluster_name, service_name)
    old_task_count = service['desiredCount']

    logger.info(
        f'Scaling service {service_name} in cluster {cluster_name}. '
        f'Current desired task count: {old_task_count}. '
        f'Current running count: {service["runningCount"]}. '
        f'New task count: {desired_count}')

    ecs_client.update_service(
        cluster=cluster_name,
        service=service_name,
        desiredCount=desired_count
    )

    if wait:
        logger.info('Waiting for service to be stable')
        ecs_client.get_waiter('services_stable').wait(
            cluster=cluster_name,
            services=[service_name]
        )

    updated_service = describe_service(ecs_client, cluster_name, service_name)
    logger.info(
        f'After scaling service {service_name} in cluster {cluster_name}. '
        f'running task count: {updated_service["runningCount"]}. '
        f'pending task count: {updated_service["pendingCount"]}. '
        f'desired count: {updated_service["desiredCount"]}')

    return old_task_count


def scale_in_service(ecs_client, event, wait=True):
    """Scale in an ECS service."""
    return scale_service(
        ecs_client=ecs_client,
        cluster_name=event['cluster_name'],
        service_name=event['service_name'],
        desired_count=0,
        wait=wait
    )


def scale_out_service(ecs_client, event, wait=True):
    """Scale out an ECS service."""
    task_count = event['task_count'] if event['task_count'] > 0 else 1
    return scale_service(
        ecs_client=ecs_client,
        cluster_name=event['cluster_name'],
        service_name=event['service_name'],
        desired_count=task_count,
        wait=wait
    )


def reset_kafka_topics_config(event, kafka_client):
    """Reset cleanup.policy for the given Kafka topics."""
    topics_to_delete = event.get('topics', [])
    topics = kafka_client.list_topics()
    topic_config = {'cleanup.policy': 'compact'}
    for topic in topics_to_delete:
        if topic in topics:
            # topic got re-created before task start, change config.
            logger.info(f'Changing config for Kafka topic {topic}.')
            resource = ConfigResource(
                'topic',
                topic,
                configs=topic_config
            )
            result_dict = kafka_client.alter_configs([resource])
            result = result_dict.to_object()
            logger.info(
                f'Kafka topic {topic} altered. Result: {json.dumps(result)}')
        else:
            logger.info(f'Kafka topic {topic} does not exist. Continuing.')


def reset_kafka_topics(event, kafka_client):
    """Delete the given Kafka topics."""
    topics_to_delete = event.get('topics', [])

    topics = kafka_client.list_topics()

    for topic in topics_to_delete:
        if topic in topics:
            logger.info(f'Deleting Kafka topic {topic}.')
            result = kafka_client.delete_topics(topics=[topic])
            res_obj = result.to_object()
            logger.info(f'Kafka topic {topic} deleted. Result: {json.dumps(res_obj)}')  # noqa
        else:
            logger.info(f'Kafka topic {topic} does not exist. Continuing.')


def delete_connector(event):
    """Make Kafka DELETE api call for the debezium connector."""
    service_name = event.get('service_name')
    connector_name = event.get('connector_name', 'debezium_mysql_source')
    if not service_name:
        return
    try:
        res = requests.delete(
            f'https://{service_name}.theorchard.io/connectors/{connector_name}')  # noqa
        logger.info(
            f'Debezium connector deleted.Result: {res.status_code} {res.text}')
    except requests.exceptions.RequestException as err:
        logger.info(f'Failed to delete Debezium connector. Error: {err}')


def handler(event, _context):
    """Lambda entry point."""
    try:
        credentials = assume_target_account_role(
            event['account_id'],
            event['role']
        )
        ecs_client = boto3.client(
            'ecs',
            region_name=config.AWS_DEFAULT_REGION,
            aws_access_key_id=credentials['AccessKeyId'],
            aws_secret_access_key=credentials['SecretAccessKey'],
            aws_session_token=credentials['SessionToken']
        )

        scale_action = event.get('scale_action')
        wait = event.get('wait', True)

        if scale_action == 'in':
            delete_connector(event)
            original_task_count = scale_in_service(
                ecs_client, event, wait=wait)
            kafka_client = KafkaAdminClient(
                bootstrap_servers=event['bootstrap_servers'],
                security_protocol='SSL'
            )
            reset_kafka_topics(event, kafka_client)
            return {
                'original_task_count': original_task_count
            }
        elif scale_action == 'out':
            kafka_client = KafkaAdminClient(
                bootstrap_servers=event['bootstrap_servers'],
                security_protocol='SSL'
            )
            reset_kafka_topics_config(event, kafka_client)
            # this will create the connector that was deleted during scale-in
            original_task_count = scale_out_service(
                ecs_client, event, wait=wait)
        else:
            raise Exception(f'Unsupported scale action {scale_action}')

        return {
            'original_task_count': original_task_count
        }
    except Exception as error:
        logger.exception(str(error))
        raise error