"""
Author: Hexu
Date: 2022-05-05 12:07:16
LastEditors: Hexu
LastEditTime: 2024/9/23 09:56
FileName: __init__.py
Description:
"""
import os

from intelliw.utils.kafka.client import KafkaClient, handle_message, delivery_callback

from intelliw.utils.logger import _get_framework_logger

logger = _get_framework_logger()


class KafkaAuthTypeEnum:
    PLAIN = 'plain'
    SCRAM = 'scram'
    SCRAM_SSL = 'scram_ssl'
    SSL = 'ssl'


AUTH_TYPE_CONFIG = {
    KafkaAuthTypeEnum.PLAIN: {
        'security_protocol': 'SASL_PLAINTEXT',
        'sasl_mechanism': 'PLAIN'
    },
    KafkaAuthTypeEnum.SCRAM: {
        'security_protocol': 'SASL_PLAINTEXT',
        'sasl_mechanism': 'SCRAM-SHA-256'
    },
    KafkaAuthTypeEnum.SCRAM_SSL: {
        'security_protocol': 'SASL_SSL',
        'sasl_mechanism': 'SCRAM-SHA-256'
    },
    KafkaAuthTypeEnum.SSL: {
        'security_protocol': 'SSL',
        'sasl_mechanism': None
    },
}


def _get_args(brokers: str, username: str, password: str, auth_type: str = None):
    """
    获取 Kafka 连接参数，优先使用传入参数，若未提供则从环境变量中获取
    """
    args = {
        'brokers': brokers or os.environ.get('kafka.cluster.hosts', ''),
    }

    if not args['brokers']:
        raise KeyError('Brokers information is missing')

    kafka_auth_type = auth_type or os.environ.get('kafka.authentication.type', '')

    if kafka_auth_type:
        auth_type_lower = kafka_auth_type.lower()
        if auth_type_lower in AUTH_TYPE_CONFIG:
            config = AUTH_TYPE_CONFIG[auth_type_lower]
            args['security_protocol'] = config['security_protocol']
            args['sasl_mechanism'] = config['sasl_mechanism']
            args['auth_type'] = kafka_auth_type
        else:
            logger.warning(f"Unknown kafka.authentication.type: {kafka_auth_type}, will use default")
    else:
        kafka_auth_enable = os.environ.get('kafka.authentication.enable', '')
        if kafka_auth_enable.lower() == 'true':
            args['username'] = username or os.environ.get('kafka.username', '')
            args['password'] = password or os.environ.get('kafka.password', '')

    if kafka_auth_type in (KafkaAuthTypeEnum.PLAIN, KafkaAuthTypeEnum.SCRAM, KafkaAuthTypeEnum.SCRAM_SSL):
        args['username'] = username or os.environ.get('kafka.username', '')
        args['password'] = password or os.environ.get('kafka.password', '')
    elif not kafka_auth_type:
        kafka_auth_enable = os.environ.get('kafka.authentication.enable', '')
        if kafka_auth_enable.lower() == 'true':
            args['username'] = username or os.environ.get('kafka.username', '')
            args['password'] = password or os.environ.get('kafka.password', '')
            args['security_protocol'] = 'SASL_PLAINTEXT'
            args['sasl_mechanism'] = 'PLAIN'

    return args


def get_client(brokers: str = None, group_id: str = None, username: str = None, password: str = None,
               producer_config: dict = None, consumer_config: dict = None,
               security_protocol: str = None, sasl_mechanism: str = None,
               auth_type: str = None,
               key_serializer=None, value_serializer=None, ssl_ca_location: str = None,
               ssl_endpoint_identification_algorithm: str = 'none'):
    """
    创建 Kafka 客户端实例
    :param brokers: Kafka 集群地址
    :param group_id: 消费者分组 ID
    :param username: SASL 验证用的用户名
    :param password: SASL 验证用的密码
    :param producer_config: 生产者配置
    :param consumer_config: 消费者配置
    :param security_protocol: 安全协议，若设置 auth_type 则自动推断
    :param sasl_mechanism: SASL 认证机制，若设置 auth_type 则自动推断
    :param auth_type: 认证类型 (plain, scram, scram_ssl, ssl)
    :param key_serializer: 用于序列化 key 的函数
    :param value_serializer: 用于序列化 value 的函数
    :param ssl_ca_location: SSL 证书路径，若不指定则默认为 /etc/ssl/certs/ca-certificates.crt
    :param ssl_endpoint_identification_algorithm: SSL 端点标识算法，默认为 'none'，可设为 'hostname' 启用主机名验证
    :return: KafkaClient 实例或 None（如果获取参数失败）
    """

    try:
        args = _get_args(brokers, username, password, auth_type)

        if ssl_ca_location is None:
            ssl_ca_location = os.environ.get('kafka.ssl.ca.location', '/etc/ssl/certs/ca-certificates.crt')

        kafka_client = KafkaClient(
            brokers=args['brokers'],
            group_id=group_id,
            username=args.get('username'),
            password=args.get('password'),
            producer_config=producer_config,
            consumer_config=consumer_config,
            security_protocol=security_protocol or args.get('security_protocol'),
            sasl_mechanism=sasl_mechanism or args.get('sasl_mechanism'),
            auth_type=auth_type or args.get('auth_type'),
            key_serializer=key_serializer,
            value_serializer=value_serializer,
            ssl_ca_location=ssl_ca_location,
            ssl_endpoint_identification_algorithm=ssl_endpoint_identification_algorithm
        )
        return kafka_client

    except KeyError as e:
        logger.error(f"Kafka args error: {e}")
        return None
