#!/usr/bin/env python3
"""
Kafka Client V2 - 基于 kafka-python 的实现
与原始 client.py 保持接口兼容，但使用 kafka-python 库
"""
import gc
import time
import json
import traceback
from concurrent.futures import ThreadPoolExecutor

try:
    from kafka import KafkaProducer, KafkaConsumer
    from kafka.errors import KafkaError
except ImportError:
    raise ImportError("\033[31mIf use kafka-python, you need: pip install kafka-python>=2.0.2\033[0m")

import threading
from intelliw.utils.logger import _get_framework_logger

logger = _get_framework_logger()


class KafkaClient:
    """
    Kafka 通用客户端 - 基于 kafka-python 实现
    与原始 client.py 保持完全接口兼容
    """

    def __init__(self, brokers: str, 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 集群的地址 (例如 'localhost:9092')
        :param group_id: 消费者使用的分组ID，生产者可忽略
        :param username: SASL 验证用的用户名
        :param password: SASL 验证用的密码
        :param producer_config: 自定义的生产者配置字典
        :param consumer_config: 自定义的消费者配置字典
        :param security_protocol: 安全协议，例如 'SASL_SSL' 或 'SASL_PLAINTEXT'
        :param sasl_mechanism: SASL 认证机制，'PLAIN' 适用于用户名和密码的认证方式
        :param key_serializer: 用于序列化 key 的函数
        :param value_serializer: 用于序列化 value 的函数
        """
        self.brokers = brokers
        self.group_id = group_id
        self.username = username
        self.password = password
        self.security_protocol = security_protocol
        self.sasl_mechanism = sasl_mechanism
        self.producer_config = producer_config or {}
        self.consumer_config = consumer_config or {}
        self.key_serializer = key_serializer or self.default_serializer
        self.value_serializer = value_serializer or self.default_serializer

        self.producer = None

        logger.info(
            f"Kafka Client(kafka-python) initialized with brokers: {brokers}, "
            f"producer_config: {self.producer_config}, "
            f"consumer_config: {self.consumer_config}")

    @staticmethod
    def default_serializer(data):
        """
        默认序列化器
        与原始 client.py 的 default_serializer 保持相同行为
        """
        if isinstance(data, str):
            return data.encode('utf-8')
        elif isinstance(data, (dict, list, tuple)):
            return json.dumps(data, ensure_ascii=False).encode('utf-8')
        elif isinstance(data, (int, float, bool)):
            return str(data).encode('utf-8')
        return data

    def _get_security_config(self):
        """
        获取 SASL 安全配置
        与原始 client.py 的 _get_security_config 保持相同配置
        """
        config = {}
        if self.username and self.password:
            config.update({
                'security_protocol': self.security_protocol,
                'sasl_mechanism': self.sasl_mechanism,
                'sasl_plain_username': self.username,
                'sasl_plain_password': self.password,
            })
        return config

    def create_producer(self):
        """
        创建 Kafka 生产者，优先使用用户传入配置
        与原始 client.py 的 create_producer 保持相同默认配置
        """
        default_producer_conf = {
            'bootstrap_servers': self.brokers,
            'batch_size': 1000,  # 批量发送的最大消息数，对应 batch.num.messages
            'linger_ms': 10,  # 批量发送时等待时间，对应 linger.ms
            'buffer_memory': 33554432,  # 缓冲区大小，对应 queue.buffering.max.kbytes
            'retries': 5,  # 消息发送失败重试次数
            'acks': 'all',  # 消息确认级别
            'compression_type': 'gzip',  # 压缩类型
            'request_timeout_ms': 30000,  # 消息超时配置，对应 message.timeout.ms
        }
        
        security_conf = self._get_security_config()
        config = {**default_producer_conf, **security_conf, **self.producer_config}
        
        # 移除 kafka-python 不支持的配置
        config.pop('queue.buffering.max.ms', None)
        config.pop('message.timeout.ms', None)
        
        return KafkaProducer(**config)

    def create_consumer(self, topics: list):
        """
        创建 Kafka 消费者，优先使用用户传入配置
        与原始 client.py 的 create_consumer 保持相同默认配置
        """
        default_consumer_conf = {
            'bootstrap_servers': self.brokers,
            'group_id': self.group_id,
            'auto_offset_reset': 'earliest',  # 从最早的偏移量开始消费
            'enable_auto_commit': False,  # 手动提交偏移量
            'max_poll_records': 500,  # 每次拉取的最大记录数
        }
        
        # 转换配置参数格式，兼容 confluent-kafka 的点分隔格式
        converted_config = {}
        for key, value in self.consumer_config.items():
            # 将点分隔格式转换为下划线格式（kafka-python格式）
            converted_key = key.replace('.', '_')
            converted_config[converted_key] = value
        
        security_conf = self._get_security_config()  # 获取 SASL 安全配置
        config = {**default_consumer_conf, **security_conf, **converted_config}
        
        # 过滤掉不支持的配置参数
        supported_configs = {
            'bootstrap_servers', 'client_id', 'group_id', 'key_deserializer',
            'value_deserializer', 'fetch_min_bytes', 'fetch_max_wait_ms',
            'max_partition_fetch_bytes', 'send_buffer_bytes', 'receive_buffer_bytes',
            'consumer_timeout_ms', 'auto_offset_reset', 'enable_auto_commit',
            'max_poll_records', 'max_poll_interval_ms', 'session_timeout_ms',
            'heartbeat_interval_ms', 'partition_assignment_strategy',
            'metric_reporters', 'metrics_num_samples', 'metrics_sample_window_ms',
            'metric_group_prefix', 'api_version_auto_timeout_ms', 'connections_max_idle_ms',
            'request_timeout_ms', 'retry_backoff_ms', 'reconnect_backoff_ms',
            'reconnect_backoff_max_ms', 'max_in_flight_requests_per_connection',
            'security_protocol', 'ssl_context', 'ssl_check_hostname',
            'ssl_cafile', 'ssl_certfile', 'ssl_keyfile', 'ssl_password',
            'ssl_crlfile', 'ssl_ciphers', 'api_version', 'selector',
            'exclude_internal_topics', 'sasl_mechanism', 'sasl_plain_username',
            'sasl_plain_password', 'sasl_kerberos_service_name',
            'sasl_kerberos_domain_name', 'sasl_oauth_token_provider'
        }
        
        final_config = {}
        for key, value in config.items():
            if key in supported_configs:
                final_config[key] = value
            else:
                # 记录被过滤的配置项，但不中断程序
                import logging
                logging.getLogger(__name__).warning(f"跳过不支持的Kafka配置参数: {key}")
        
        consumer = KafkaConsumer(**final_config)
        consumer.subscribe(topics)
        return KafkaConsumerWrapper(consumer)

    def produce_message(self, topic: str, key: str = None, value: str = None, raise_error=False):
        """
        发送 Kafka 消息
        与原始 client.py 的 produce_message 保持相同接口和行为
        """
        assert not (value is None and key is None), 'Need at least one: key or value'

        serialized_key = self.key_serializer(key)
        serialized_value = self.value_serializer(value)

        # 确保序列化后的数据类型正确
        if serialized_key is not None:
            assert type(serialized_key) in (bytes, bytearray, memoryview, type(None))
        if serialized_value is not None:
            assert type(serialized_value) in (bytes, bytearray, memoryview, type(None))

        if self.producer is None:
            self.producer = self.create_producer()

        try:
            future = self.producer.send(topic, key=serialized_key, value=serialized_value)
            if raise_error:
                # 同步等待结果并抛出异常
                result = future.get(timeout=30)
                logger.info(f"Message produced to topic '{topic}' successfully")
                return result is not None
            else:
                # 异步回调方式
                future.add_callback(lambda m: delivery_callback(None, m))
                future.add_errback(lambda e: delivery_callback(e, None))
                self.producer.flush()
                return True
        except Exception as e:
            logger.error(f"Error producing message to topic '{topic}', key: {serialized_key}, error: {e}")
            if raise_error:
                raise  # 直接抛出原异常，保留异常信息
            return False
        finally:
            self.producer.flush()

    def produce_messages_in_batch(self, topic: str, messages: list, poll_interval: float = 0.1):
        """
        批量发送 Kafka 消息
        与原始 client.py 的 produce_messages_in_batch 保持相同接口和行为
        :param topic: Kafka topic
        :param messages: 消息列表，包含 (key, value) 元组
        :param poll_interval: 每隔多少秒调用一次 poll() 以确保消息发送
        """
        if self.producer is None:
            self.producer = self.create_producer()

        try:
            futures = []
            for key, value in messages:
                serialized_key = self.key_serializer(key)
                serialized_value = self.value_serializer(value)
                future = self.producer.send(topic, key=serialized_key, value=serialized_value)
                futures.append(future)
                self.producer.flush()  # 定期调用以确保消息发送

            # 等待所有消息发送完成
            while futures:
                remaining_futures = []
                for future in futures:
                    if not future.is_done:
                        remaining_futures.append(future)
                
                if not remaining_futures:
                    break
                
                time.sleep(poll_interval)
                self.producer.flush()
            
            self.producer.flush()
            logger.info(f"Successfully produced {len(messages)} messages to topic '{topic}'")
            return True
        except Exception as e:
            logger.error(f"Error producing message batch: {e}")
            return False

    def consume_messages(self, topics: list, handle_func, timeout=1.0, return_by_empty=False, return_by_count=0,
                         return_by_func=None):
        """
        消费 Kafka 消息
        与原始 client.py 的 consume_messages 保持相同接口和行为
        :param topics: Kafka topic list
        :param handle_func: message 后处理方法，
            方法需要有两个参数，例如：
             - def handle_message(key, value)
             - def delivery_callback(err, msg)
            第一个参数代表consumer的 message.key(), 第二个参数为consumer的 message.value()
        :param timeout: 获取消息超时时长
        :param return_by_empty: 当获取消息 None 时结束
        :param return_by_count 当获取消息数量等于count时结束, 0为不限制
        :param return_by_func  参数传入一个方法，方法参数与handle_func相同，返回类型需要为 bool 型，当返回为 True 时，消费停止
        """
        consumer = self.create_consumer(topics)

        try:
            count = 0
            while True:
                messages = consumer.poll(timeout_ms=int(timeout * 1000))
                
                if not messages:
                    # 当获取消息 None 时结束
                    if return_by_empty:
                        return None
                    continue

                for topic_partition, records in messages.items():
                    for record in records:
                        try:
                            # 解码消息key和value
                            key = record.key.decode('utf-8') if record.key else None
                            value = record.value.decode('utf-8') if record.value else None
                            
                            # 尝试解析 JSON，与原始 client.py 保持一致
                            try:
                                if value:
                                    value = json.loads(value)
                                if key:
                                    key = json.loads(key)
                            except (json.JSONDecodeError, TypeError):
                                pass
                            
                            # 处理消息
                            handle_func(key, value)  # 处理消息
                            consumer.commit()  # 手动提交偏移量
                            count += 1

                            # 当获取消息数量等于count时结束
                            if return_by_count and count >= return_by_count:
                                return None
                            if return_by_func and return_by_func(key, value) is True:
                                return None

                            # 定期调用垃圾回收
                            if count % 10 == 0:
                                gc.collect()
                        except Exception as e:
                            logger.error(f"Error consuming messages: {e}, {traceback.format_exc()}")

        except Exception as e:
            logger.error(f"Error consuming messages: {e}, {traceback.format_exc()}")
        finally:
            consumer.close()

        return None

    def consume_messages_with_threads(self, topics: list, handle_func, num_threads=5, timeout=1.0):
        """
        多线程消费 Kafka 消息
        与原始 client.py 的 consume_messages_with_threads 保持相同接口和行为
        """
        def consumer_task():
            """单个消费者线程任务"""
            self.consume_messages(topics, handle_func, timeout)

        with ThreadPoolExecutor(max_workers=num_threads) as executor:
            futures = [executor.submit(consumer_task) for _ in range(num_threads)]
            for future in futures:
                future.result()  # 等待所有线程完成任务


def handle_message(key, value):
    """
    默认消息处理函数
    与原始 client.py 的 handle_message 保持相同日志格式
    """
    logger.info(f"Received message with key: {key}, value: {value}")
    print(f"Received message with key: {key}, value: {value}")


class KafkaConsumerWrapper:
    """
    KafkaConsumer 包装类，提供与 confluent-kafka Consumer 兼容的 API
    """
    def __init__(self, consumer):
        self._consumer = consumer
    
    def poll(self, timeout=None, max_records=None, **kwargs):
        """
        兼容 confluent-kafka 的 poll 方法
        
        :param timeout: 超时时间（秒），与 confluent-kafka 保持一致
        :param max_records: 最大记录数，可选参数
        :return: 消息或 None
        """
        # 处理不同的参数格式
        timeout_ms = None
        
        # 优先处理 timeout 参数（confluent-kafka 格式）
        if timeout is not None:
            timeout_ms = int(timeout * 1000)  # 转换为毫秒
        elif 'timeout' in kwargs:
            timeout_ms = int(kwargs.pop('timeout') * 1000)
        elif 'timeout_ms' in kwargs:
            timeout_ms = kwargs.pop('timeout_ms')
            
        try:
            # 调用 kafka-python 的 poll 方法
            messages = self._consumer.poll(timeout_ms=timeout_ms, max_records=max_records, **kwargs)
            
            # 将 kafka-python 的返回格式转换为与 confluent-kafka 兼容的格式
            if not messages:
                return None
                
            # 返回第一条消息（与 confluent-kafka 行为一致）
            for topic_partition, records in messages.items():
                if records:
                    record = records[0]
                    # 创建兼容的消息对象
                    return KafkaMessageWrapper(record)
                    
            return None
            
        except Exception as e:
            logger.error(f"Error in poll: {e}")
            return None
    
    def close(self):
        """关闭消费者"""
        return self._consumer.close()
    
    def subscribe(self, topics):
        """订阅主题"""
        return self._consumer.subscribe(topics)
    
    def unsubscribe(self):
        """取消订阅"""
        return self._consumer.unsubscribe()
    
    def assign(self, partitions):
        """分配分区"""
        return self._consumer.assign(partitions)
    
    def assignment(self):
        """获取当前分配的分区"""
        return self._consumer.assignment()
    
    def commit(self, offsets=None):
        """提交偏移量"""
        return self._consumer.commit(offsets)
    
    def committed(self, partition):
        """获取已提交的偏移量"""
        return self._consumer.committed(partition)
    
    def position(self, partition):
        """获取当前位置"""
        return self._consumer.position(partition)
    
    def seek(self, partition, offset):
        """设置位置"""
        return self._consumer.seek(partition, offset)
    
    def seek_to_beginning(self, *partitions):
        """移动到开头"""
        return self._consumer.seek_to_beginning(*partitions)
    
    def seek_to_end(self, *partitions):
        """移动到结尾"""
        return self._consumer.seek_to_end(*partitions)
    
    def partitions_for_topic(self, topic):
        """获取主题的分区"""
        return self._consumer.partitions_for_topic(topic)
    
    def metrics(self, raw=False):
        """获取指标"""
        return self._consumer.metrics(raw=raw)
    
    def __getattr__(self, name):
        """委托其他属性到内部消费者"""
        return getattr(self._consumer, name)


class KafkaMessageWrapper:
    """
    Kafka 消息包装类，提供与 confluent-kafka Message 兼容的 API
    """
    def __init__(self, record):
        self._record = record
        self.topic = lambda: record.topic
        self.partition = lambda: record.partition
        self.offset = lambda: record.offset
        self.timestamp = lambda: record.timestamp
        self.key = lambda: record.key
        self.value = lambda: record.value
        self.headers = lambda: record.headers if hasattr(record, 'headers') else None
        self.error = lambda: None  # kafka-python 没有错误概念，总是返回 None
        self.topic.return_value = record.topic
        self.partition.return_value = record.partition
        self.offset.return_value = record.offset
        self.timestamp.return_value = record.timestamp
        self.key.return_value = record.key
        self.value.return_value = record.value
        self.headers.return_value = record.headers if hasattr(record, 'headers') else None
        self.error.return_value = None


def delivery_callback(err, msg):
    """
    消息发送回调函数
    与原始 client.py 的 delivery_callback 保持相同日志格式
    """
    if err:
        logger.error(f"Message failed delivery: {err}")
        print(f"Message failed delivery: {err}")
    else:
        logger.info(f"Message delivered to {msg.topic} [{msg.partition}] at offset {msg.offset}")