"""
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.logger import _get_framework_logger

logger = _get_framework_logger()

# 检查是否安装了 confluent_kafka 或 kafka-python

debug_kafka_python = os.environ.get('DEBUG_KAFKA_PYTHON', 'false').lower() == 'true'
print("debug_kafka_python:", debug_kafka_python)
has_kafka_package = False

# 根据 debug_kafka_python 决定导入顺序
if debug_kafka_python:
    # debug_kafka_python=true 时，优先使用 client_v2 (kafka-python)
    try:
        from intelliw.utils.kafka.client_v2 import KafkaClient, handle_message, delivery_callback
        logger.info("kafka-python (client_v2) is being used")
        has_kafka_package = True
    except (ImportError, ModuleNotFoundError):
        logger.warning("kafka-python (client_v2) not installed, trying confluent_kafka (client)")
        try:
            from intelliw.utils.kafka.client import KafkaClient, handle_message, delivery_callback
            logger.info("confluent_kafka (client) is being used")
            has_kafka_package = True
        except (ImportError, ModuleNotFoundError):
            pass
else:
    # debug_kafka_python=false 时，优先使用 client (confluent_kafka)
    try:
        from intelliw.utils.kafka.client import KafkaClient, handle_message, delivery_callback
        logger.info("confluent_kafka (client) is being used")
        has_kafka_package = True
    except (ImportError, ModuleNotFoundError):
        logger.warning("confluent_kafka (client) not installed, trying kafka-python (client_v2)")
        try:
            from intelliw.utils.kafka.client_v2 import KafkaClient, handle_message, delivery_callback
            logger.info("kafka-python (client_v2) is being used")
            has_kafka_package = True
        except (ImportError, ModuleNotFoundError):
            pass

if not has_kafka_package:
    raise ImportError("\033[31mIf use kafka, you need: pip install confluent_kafka>=2.5.3 or pip install kafka-python\033[0m")



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

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

    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', '')

    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 = 'SASL_PLAINTEXT', sasl_mechanism: str = 'PLAIN',
               key_serializer=None, value_serializer=None):
    """
    创建 Kafka 客户端实例
    :param brokers: Kafka 集群地址
    :param group_id: 消费者分组 ID
    :param username: SASL 验证用的用户名
    :param password: SASL 验证用的密码
    :param producer_config: 生产者配置
    :param consumer_config: 消费者配置
    :param security_protocol: 安全协议，默认为 SASL_PLAINTEXT
    :param sasl_mechanism: SASL 认证机制，默认为 PLAIN
    :param key_serializer: 用于序列化 key 的函数
    :param value_serializer: 用于序列化 value 的函数
    :return: KafkaClient 实例或 None（如果获取参数失败）
    """

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

        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,
            sasl_mechanism=sasl_mechanism,
            key_serializer=key_serializer,
            value_serializer=value_serializer
        )
        return kafka_client

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