"""
Kafka queue implementation
Requires kafka-python and protobuf libraries to be installed
Supports both Protocol Buffers and JSON serialization
"""
import json
import logging
from typing import Optional

from google.protobuf.json_format import MessageToDict, ParseDict
from mirix.queue.queue_interface import QueueInterface
from mirix.queue.message_pb2 import QueueMessage

logger = logging.getLogger(__name__)


class KafkaQueue(QueueInterface):
    """Kafka-based queue implementation supporting Protobuf and JSON serialization"""
    
    def __init__(
        self, 
        bootstrap_servers: str, 
        topic: str, 
        group_id: str,
        serialization_format: str = 'protobuf',
        security_protocol: str = 'PLAINTEXT',
        ssl_cafile: Optional[str] = None,
        ssl_certfile: Optional[str] = None,
        ssl_keyfile: Optional[str] = None
    ):
        """
        Initialize Kafka producer and consumer with configurable serialization
        
        Args:
            bootstrap_servers: Kafka broker address(es)
            topic: Kafka topic name
            group_id: Consumer group ID
            serialization_format: 'protobuf' or 'json' (default: 'protobuf')
            security_protocol: Kafka security protocol - 'PLAINTEXT', 'SSL', 'SASL_PLAINTEXT', 'SASL_SSL'
            ssl_cafile: Path to CA certificate file for SSL/TLS verification
            ssl_certfile: Path to client certificate file for mTLS
            ssl_keyfile: Path to client private key file for mTLS
        """
        logger.info(
            "🔧 Initializing Kafka queue: servers=%s, topic=%s, group=%s, format=%s, security=%s", 
            bootstrap_servers, topic, group_id, serialization_format, security_protocol
        )
        
        try:
            from kafka import KafkaProducer, KafkaConsumer
        except ImportError:
            logger.error("kafka-python not installed")
            raise ImportError(
                "kafka-python is required for Kafka support. "
                "Install it with: pip install queue-sample[kafka]"
            )
        
        self.topic = topic
        self.serialization_format = serialization_format.lower()
        
        # Protobuf serializer: Convert QueueMessage to bytes
        def protobuf_serializer(message: QueueMessage) -> bytes:
            """
            Serialize QueueMessage to Protocol Buffer format
            
            Args:
                message: QueueMessage protobuf to serialize
                
            Returns:
                Serialized protobuf bytes
            """
            return message.SerializeToString()
        
        # Protobuf deserializer: Convert bytes to QueueMessage
        def protobuf_deserializer(serialized_msg: bytes) -> QueueMessage:
            """
            Deserialize Protocol Buffer message to QueueMessage
            
            Args:
                serialized_msg: Serialized protobuf bytes
                
            Returns:
                QueueMessage protobuf object
            """
            msg = QueueMessage()
            msg.ParseFromString(serialized_msg)
            return msg
        
        # JSON serializer: Convert QueueMessage to JSON bytes
        def json_serializer(message: QueueMessage) -> bytes:
            """
            Serialize QueueMessage to JSON format
            
            Args:
                message: QueueMessage protobuf to serialize
                
            Returns:
                JSON bytes
            """
            message_dict = MessageToDict(message, preserving_proto_field_name=True)
            return json.dumps(message_dict).encode('utf-8')
        
        # JSON deserializer: Convert JSON bytes to QueueMessage
        def json_deserializer(serialized_msg: bytes) -> QueueMessage:
            """
            Deserialize JSON message to QueueMessage
            
            Args:
                serialized_msg: JSON bytes
                
            Returns:
                QueueMessage protobuf object
            """
            message_dict = json.loads(serialized_msg.decode('utf-8'))
            return ParseDict(message_dict, QueueMessage())
        
        # Select serializer/deserializer based on format
        if self.serialization_format == 'json':
            value_serializer = json_serializer
            value_deserializer = json_deserializer
            logger.info("Using JSON serialization for Kafka messages")
        else:
            value_serializer = protobuf_serializer
            value_deserializer = protobuf_deserializer
            logger.info("Using Protobuf serialization for Kafka messages")
        
        # Build Kafka producer/consumer config with optional SSL
        kafka_config = {
            'bootstrap_servers': bootstrap_servers,
        }
        
        # Add SSL configuration if security protocol is SSL
        if security_protocol.upper() in ['SSL', 'SASL_SSL']:
            kafka_config['security_protocol'] = security_protocol.upper()
            if ssl_cafile:
                kafka_config['ssl_cafile'] = ssl_cafile
            if ssl_certfile:
                kafka_config['ssl_certfile'] = ssl_certfile
            if ssl_keyfile:
                kafka_config['ssl_keyfile'] = ssl_keyfile
            logger.info("Kafka SSL/TLS configured: protocol=%s", security_protocol)
        
        # Initialize Kafka producer with selected serializer and key serializer
        # Key serializer enables partition key routing for consistent message ordering per user
        self.producer = KafkaProducer(
            **kafka_config,
            key_serializer=lambda k: k.encode('utf-8'),  # Encode partition key to bytes
            value_serializer=value_serializer
        )
        
        # Initialize Kafka consumer with selected deserializer
        self.consumer = KafkaConsumer(
            topic,
            **kafka_config,
            group_id=group_id,
            value_deserializer=value_deserializer,
            auto_offset_reset='earliest',  # Start from beginning if no offset exists
            enable_auto_commit=True,
            consumer_timeout_ms=1000  # Timeout for polling
        )
    
    def put(self, message: QueueMessage) -> None:
        """
        Send a message to Kafka topic with user_id as partition key.
        
        This ensures all messages for the same user go to the same partition,
        guaranteeing single-worker processing and message ordering per user.
        
        Implementation:
        - Uses user_id (or actor.id as fallback) as partition key
        - Kafka assigns partition via: hash(key) % num_partitions
        - Consumer group ensures only one worker per partition
        - Result: Same user always processed by same worker (no race conditions)
        
        Args:
            message: QueueMessage protobuf message to send
        """
        # Extract user_id as partition key (fallback to actor.id if not present)
        partition_key = message.user_id if message.user_id else message.actor.id
        
        logger.debug(
            "Sending message to Kafka topic %s: agent_id=%s, partition_key=%s",
            self.topic, message.agent_id, partition_key
        )
        
        # Send message with partition key - ensures consistent partitioning
        # Kafka will route this to: partition = hash(partition_key) % num_partitions
        future = self.producer.send(
            self.topic,
            key=partition_key,  # Partition key for consistent routing
            value=message
        )
        future.get(timeout=10)  # Wait up to 10 seconds for confirmation
        
        logger.debug("Message sent to Kafka successfully with partition key: %s", partition_key)
    
    def get(self, timeout: Optional[float] = None) -> QueueMessage:
        """
        Retrieve a message from Kafka
        
        Args:
            timeout: Not used for Kafka (uses consumer_timeout_ms instead)
            
        Returns:
            QueueMessage protobuf message from Kafka
            
        Raises:
            StopIteration: If no message available
        """
        logger.debug("Polling Kafka topic %s for messages", self.topic)
        
        # Poll for messages
        for message in self.consumer:
            logger.debug("Retrieved message from Kafka: agent_id=%s", message.value.agent_id)
            return message.value
        
        # If no message received, raise exception (similar to queue.Empty)
        logger.debug("No message available from Kafka")
        raise StopIteration("No message available")
    
    def close(self) -> None:
        """Close Kafka producer and consumer connections"""
        logger.info("Closing Kafka connections")
        
        if hasattr(self, 'producer'):
            self.producer.close()
            logger.debug("Kafka producer closed")
        if hasattr(self, 'consumer'):
            self.consumer.close()
            logger.debug("Kafka consumer closed")

