"""AWS helper functions."""
from time import sleep
from typing import Any
from typing import Dict
from typing import Optional
from typing import Union

import boto3
from botocore.client import BaseClient
from botocore.credentials import Credentials
import requests
from secrets_manager.python_ext import PythonSecretsManager

from dbdeploy.base import config


def get_aws_creds() -> Credentials:
    """Get AWS credentials."""
    return boto3.Session().get_credentials()


def get_secret_name(
        environment: str,
        service_name: str,
        secret_key: str) -> str:
    """Return secret name."""
    return '{environment}/{service_name}/{secret_key}'.format(
        environment=environment,
        service_name=service_name,
        secret_key=secret_key)


def get_secret(
        aws_region: str,
        environment: str,
        service_name: str,
        secret_key: str) -> Dict[str, str]:
    """Get credential dict from AWS SecretsManager."""
    secret_name = get_secret_name(environment, service_name, secret_key)

    dict_secret = PythonSecretsManager(
        region_name=aws_region,
        environment=environment,
    ).get_secret(secret_name=secret_name)

    if not isinstance(dict_secret, dict):
        raise ValueError('Got secret from SecretsManager, '
                         'but the value is not a JSON string.')
    if not dict_secret:
        raise ValueError('AWS SecretsManager secret is empty.')

    return dict_secret


def get_session_for_diff_aws_roles(role: str) -> boto3.Session:
    """Get aws session object assuming a different role."""
    client = boto3.client("sts")
    response = client.assume_role(
        RoleArn=role,
        RoleSessionName='kafka-db-deploy',
    )
    return boto3.Session(
        aws_access_key_id=response['Credentials']['AccessKeyId'],
        aws_secret_access_key=response['Credentials']['SecretAccessKey'],
        aws_session_token=response['Credentials']['SessionToken'],
    )


def get_cluster_info(client: boto3.Session, cluster_name: str) -> Any | None:
    """Get bootstrap brokers info."""
    cluster_list = client.list_clusters()['ClusterInfoList']
    cluster_info = None
    for cluster in cluster_list:
        if cluster_name == cluster['ClusterName']:
            cluster_info = client.get_bootstrap_brokers(
                ClusterArn=cluster['ClusterArn'])
    return cluster_info


def get_msk_bootstrap_brokers(
        cluster_name: str,
        aws_region: str = 'us-east-1') -> str:
    """Get kafka bootstrap broker servers list for provided cluster."""
    client = boto3.client('kafka', region_name=aws_region)
    cluster_info = get_cluster_info(client, cluster_name)
    if cluster_info:
        return str(cluster_info['BootstrapBrokerStringTls'])

    config.LOGGER.info('Cluster info not found on default aws account.')
    # look for this MSK cluster on other aws accounts.
    for role in config.OTHER_ROLES_FOR_CLUSTER_ACCESS:
        config.LOGGER.info(f'Checking other role {role}')
        session = get_session_for_diff_aws_roles(role)
        client = session.client('kafka', region_name=aws_region)
        cluster_info = get_cluster_info(client, cluster_name)
        if cluster_info:
            return str(cluster_info['BootstrapBrokerStringTls'])

    # cluster not found on any aws account.
    raise ValueError('Was not able to find provided kafka cluster.')


def get_ecs_client(aws_region: str = 'us-east-1') -> BaseClient:
    """Boto client factory for ECS clients.

    Returns:
        boto3 client instance
    """
    return boto3.client('ecs', region_name=aws_region)


def get_ecs_cluster_arn(
        ecs_client: BaseClient,
        cluster_name: str) -> Optional[str]:
    """Return The Amazon Resource Name (ARN) that identifies the cluster.

    Args:
        ecs_client: boto3.client('ecs') instance
        cluster_name: The short name of the ECS cluster.

    Returns:
        str: return ARN if cluster was found
    """
    cluster_description = ecs_client.describe_clusters(clusters=[cluster_name])
    clusters = cluster_description.get('clusters')
    return clusters[0]['clusterArn'] if clusters else None


def change_fargate_task_count(
        ecs_client: BaseClient,
        desired_count: int,
        cluster_name: str,
        cluster_arn: str) -> bool:
    """The number of the task instances running in the service.

    Args:
        ecs_client: boto3.client('ecs') instance
        desired_count: The desired number of running task instances
        cluster_name: The short name of the ECS cluster service.
        cluster_arn: The Amazon Resource Name (ARN) that identifies the cluster

    Returns:
        boolean: The result of task count change attempt.
    """
    ecs_client.update_service(
        desiredCount=desired_count,
        service=cluster_name,
        cluster=cluster_arn)

    return True


def check_running_task_count(
        ecs_client: BaseClient,
        desired_count: int,
        backoff_timeout: Union[int, float],
        grace_period: Union[int, float],
        cluster_arn: str,
        cluster_url: str) -> bool:
    """The number of the task instances running in the service.

    Args:
        desired_count: The desired number of running task instances
        desired_count: The desired number of running task instances
        backoff_timeout: Sleep period between checks (seconds),
        grace_period: The maximum time to check tasks state (seconds),
        cluster_arn: The Amazon Resource Name (ARN) that identifies the cluster
        cluster_url: The URL the Kafka-Connect cluster

    Returns:
        boolean: The result of change attempt.
    """
    connectors_url = f'{cluster_url}/connectors'

    while True:
        if grace_period < 0:
            return False
        task_arns = ecs_client.list_tasks(cluster=cluster_arn)['taskArns']
        task_count = int(len(task_arns))
        healthy = True
        if desired_count > 0 and task_count:
            try:
                response = requests.get(connectors_url)
                healthy = response.status_code == 200
            except Exception:
                pass
        difference = abs(task_count - desired_count)
        if healthy and not difference:
            # stop looping if all conditions are met
            break
        grace_period -= backoff_timeout
        sleep(backoff_timeout)

    return True


def get_dynamodb_client(aws_region: str = 'us-east-1') -> BaseClient:
    """Boto client factory for DynamoDB clients.

    Returns:
        boto3 client instance
    """
    return boto3.client('dynamodb', region_name=aws_region)


def get_table_item(
        dynamodb_client: BaseClient,
        changeset_id: str,
        table_name: str) -> Optional[Dict[str, str]]:
    """Get changeset status.

    Args:
        dynamodb_client: boto3.client('dynamodb') instance
        changeset_id: The id of the changeset
        table_name: The name of table

    Returns:
        Changeset status information.
    """
    response = dynamodb_client.get_item(
        TableName=table_name,
        Key={
            'changeset_id': {'S': changeset_id},
        },
        AttributesToGet=[
            'cypher_hash',
            'sql_hash'])
    item: Optional[Dict[str, str]] = response.get('Item')
    return item


def put_table_item(
        dynamodb_client: BaseClient,
        changeset_id: str,
        table_name: str,
        cypher_hash: Optional[str],
        sql_hash: Optional[str]) -> Optional[Dict[str, str]]:
    """Update changeset status.

    Args:
        dynamodb_client: boto3.client('dynamodb') instance
        changeset_id: The id of the changeset
        table_name: The name of table
        cypher_hash: Hashed changeset cypher
        sql_hash: Hashed changeset sql

    Returns:
        Updated changeset status information.
    """
    response = dynamodb_client.update_item(
        TableName=table_name,
        Key={
            'changeset_id': {'S': changeset_id},
        },
        AttributeUpdates={
            'cypher_hash': {'Value': {'S': str(cypher_hash)}},
            'sql_hash': {'Value': {'S': str(sql_hash)}},
        },
        ReturnValues='ALL_NEW')
    item: Optional[Dict[str, str]] = response.get('Attributes')
    return item


def put_table_lock_item(
        dynamodb_client: BaseClient,
        table_name: str) -> bool:
    """Set deploy lock.

    Args:
        dynamodb_client: boto3.client('dynamodb') instance
        table_name: The name of table

    Returns:
        The result
    """
    response = dynamodb_client.update_item(
        TableName=table_name,
        Key={
            'changeset_id': {'S': '__DB_DEPLOY_LOCK__'},
        },
        ReturnValues='ALL_OLD')
    return response.get('Attributes') is None


def delete_table_lock_item(
        dynamodb_client: BaseClient,
        table_name: str) -> bool:
    """Remove deploy lock.

    Args:
        dynamodb_client: boto3.client('dynamodb') instance
        table_name: The name of table

    Returns:
        The result
    """
    response = dynamodb_client.delete_item(
        TableName=table_name,
        Key={
            'changeset_id': {'S': '__DB_DEPLOY_LOCK__'},
        },
        ReturnValues='ALL_OLD')
    return response.get('Attributes') is not None