"""Kafka producer client."""

import json
import ssl

import boto3
from kafka import KafkaProducer
from lambdacommon.common_config import logger

import config


def get_switchboard_credentials():
    """Get credentials for interacting with Switchboard AWS."""
    logger.debug("Getting switchboard credentials")
    try:
        sts = boto3.client("sts")
        assumed_role = sts.assume_role(
            RoleArn=config.SWITCHBOARD_ROLE_ARN, RoleSessionName="switchboard-graphql"
        )
        credentials = assumed_role["Credentials"]
        return {
            "aws_access_key_id": credentials["AccessKeyId"],
            "aws_secret_access_key": credentials["SecretAccessKey"],
            "aws_session_token": credentials["SessionToken"],
        }
    except Exception as e:
        logger.error(f"Failed to get switchboard credentials: {str(e)}")
        raise e


def get_kafka_config(arn, credentials):
    """Get Amazon MSK broker string.

    Args:
        arn (str): The MSK Cluster's ARN
    Returns:
        Queue: Instance of Amazon SQS queue.
    """
    logger.debug("Getting bootstrap brokers")
    try:
        kafka = boto3.client("kafka", **credentials)
        broker_request = kafka.get_bootstrap_brokers(ClusterArn=arn)

        try:
            return {
                "bootstrap_servers": broker_request["BootstrapBrokerStringTls"],
                "security_protocol": "SSL",
            }
        except KeyError:
            return {
                "bootstrap_servers": broker_request["BootstrapBrokerString"],
            }
    except Exception as e:
        logger.error(f"Failed to get broker string: {str(e)}")
        raise e


def get_kafka_producer():
    """Return kafka.KafkaProducer client."""
    credentials = get_switchboard_credentials()
    kafka_config = get_kafka_config(config.KAFKA_CLUSTER_ARN, credentials)
    return KafkaProducer(
        value_serializer=lambda v: json.dumps(v).encode("utf-8"),
        ssl_context=ssl.create_default_context(),
        **kafka_config,
    )