"""
Kafka client module.

This module provides wrapper classes for Kafka producer and consumers
with async support and connection pooling.
"""

import asyncio
import json
import logging
from typing import Dict, List, Optional, Any, Callable, AsyncGenerator
from dataclasses import dataclass

try:
    from confluent_kafka import Producer, Consumer, KafkaException
    from confluent_kafka.admin import AdminClient, NewTopic
    from confluent_kafka import KafkaError as ConfluentKafkaError
except ImportError:
    raise ImportError("confluent-kafka package is required. Install it with: pip install confluent-kafka")

from kafka_config import ProducerConfig, ConsumerConfig, KafkaConfig
from kafka_errors import ProducerError, ConsumerError, TopicError, ConsumerGroupError, ConnectionError, ErrorHandler


@dataclass
class Message:
    """Kafka message wrapper."""
    topic: str
    value: str
    key: Optional[str] = None
    partition: Optional[int] = None
    headers: Optional[Dict[str, str]] = None

    def to_dict(self) -> Dict[str, Any]:
        """Convert message to dictionary."""
        return {
            "topic": self.topic,
            "value": self.value,
            "key": self.key,
            "partition": self.partition,
            "headers": self.headers
        }


class ProducerPool:
    """Producer connection pool for managing multiple producer instances."""

    def __init__(self, max_pool_size: int = 5):
        self.max_pool_size = max_pool_size
        self.producers: Dict[str, List[Producer]] = {}
        self.lock = asyncio.Lock()

    async def get_producer(self, config: ProducerConfig, name: str = 'default') -> Producer:
        """Get a producer instance from the pool or create a new one."""
        while True:
            async with self.lock:
                if name not in self.producers:
                    self.producers[name] = []

                if len(self.producers[name]) > 0:
                    return self.producers[name].pop()

                # Create new producer if pool is not full
                if len(self.producers[name]) < self.max_pool_size:
                    return self._create_producer(config)

            # Wait for a producer to be available (outside the lock)
            await asyncio.sleep(0.1)

    def _create_producer(self, config: ProducerConfig) -> Producer:
        """Create a new producer instance."""
        producer_config = {
            'bootstrap.servers': config.bootstrap_servers,
            'acks': config.acks,
            'retries': config.retries,
            'batch.size': config.batch_size,
            'linger.ms': config.linger_ms,
            'compression.type': config.compression_type,
        }

        return Producer(producer_config)

    async def return_producer(self, producer: Producer, name: str = 'default'):
        """Return a producer to the pool."""
        async with self.lock:
            if len(self.producers[name]) < self.max_pool_size:
                self.producers[name].append(producer)
            else:
                producer.close()
                producer.flush()

    def close_all(self):
        """Close all producers in the pool."""
        for name in self.producers:
            for producer in self.producers[name]:
                producer.flush()
                producer.close()


class AsyncProducer:
    """Async wrapper for Kafka producer."""

    def __init__(self, producer: Producer, pool: ProducerPool, name: str = 'default'):
        self.producer = producer
        self.pool = pool
        self.name = name
        self.logger = logging.getLogger(__name__)

    async def send_message(self, topic: str, value: str, key: Optional[str] = None,
                          headers: Optional[Dict[str, str]] = None) -> Message:
        """
        Send a single message to Kafka.

        Args:
            topic: Topic name
            value: Message value
            key: Message key (optional)
            headers: Message headers (optional)

        Returns:
            Message object with delivery information

        Raises:
            ProducerError: If message delivery fails
        """
        try:
            # Convert headers dict to list format for confluent-kafka
            headers_list = None
            if headers:
                headers_list = [(k, str(v)) for k, v in headers.items()]

            # Produce message
            future = self.producer.produce(topic, value=value, key=key, headers=headers_list)

            # Check if future is None (connection failed)
            if future is None:
                raise ProducerError(
                    f"Failed to send message to topic {topic}: producer returned None (likely connection failure)",
                    producer_id=self.name
                )

            # Wait for delivery
            delivery_result = await asyncio.get_event_loop().run_in_executor(
                None, future.get
            )

            if delivery_result.error():
                raise ProducerError(
                    f"Failed to deliver message to topic {topic}: {delivery_result.error()}",
                    producer_id=self.name
                )

            return Message(
                topic=topic,
                value=value,
                key=key,
                partition=delivery_result.partition(),
                headers=headers
            )

        except Exception as e:
            self.logger.error(f"Error sending message: {e}")
            raise ProducerError(f"Error sending message: {e}", producer_id=self.name)

    async def send_messages_batch(self, messages: List[Dict[str, Any]]) -> List[Message]:
        """
        Send multiple messages in a batch.

        Args:
            messages: List of message dictionaries

        Returns:
            List of Message objects with delivery information
        """
        results = []
        futures = []

        try:
            for message_data in messages:
                topic = message_data['topic']
                value = message_data['value']
                key = message_data.get('key')
                headers = message_data.get('headers')

                # Convert headers dict to list format for confluent-kafka
                headers_list = None
                if headers:
                    headers_list = [(k, str(v)) for k, v in headers.items()]

                future = self.producer.produce(topic, value=value, key=key, headers=headers_list)

                # Check if future is None (connection failed)
                if future is None:
                    raise ProducerError(
                        f"Failed to send message to topic {topic}: producer returned None (likely connection failure)",
                        producer_id=self.name
                    )

                futures.append((future, message_data))

            # Wait for all deliveries
            for future, message_data in futures:
                delivery_result = await asyncio.get_event_loop().run_in_executor(
                    None, future.get
                )

                if delivery_result.error():
                    raise ProducerError(
                        f"Failed to deliver message to topic {message_data['topic']}: {delivery_result.error()}",
                        producer_id=self.name
                    )

                results.append(Message(
                    topic=message_data['topic'],
                    value=message_data['value'],
                    key=message_data.get('key'),
                    partition=delivery_result.partition(),
                    headers=message_data.get('headers')
                ))

            return results

        except Exception as e:
            self.logger.error(f"Error sending messages batch: {e}")
            raise ProducerError(f"Error sending messages batch: {e}", producer_id=self.name)

    async def close(self):
        """Close the producer."""
        self.producer.flush()
        await self.pool.return_producer(self.producer, self.name)


class ConsumerManager:
    """Manager for Kafka consumers with async support."""

    def __init__(self, config: ConsumerConfig, name: str = 'default'):
        self.config = config
        self.name = name
        self.consumer = None
        self.logger = logging.getLogger(__name__)
        self.is_running = False
        self.tasks: List[asyncio.Task] = []

    def create_consumer(self) -> Consumer:
        """Create a new consumer instance."""
        consumer_config = {
            'bootstrap.servers': self.config.bootstrap_servers,
            'group.id': self.config.group_id,
            'auto.offset.reset': self.config.auto_offset_reset,
            'enable.auto.commit': self.config.enable_auto_commit,
            'auto.commit.interval.ms': self.config.auto_commit_interval_ms,
            'session.timeout.ms': self.config.session_timeout_ms,
            'heartbeat.interval.ms': self.config.heartbeat_interval_ms,
            'max.poll.records': self.config.max_poll_records,
            'max.poll.interval.ms': self.config.max_poll_interval_ms,
        }

        return Consumer(consumer_config)

    async def consume_messages(self, topics: List[str],
                              message_handler: Callable[[Dict[str, Any]], None] = None,
                              timeout: float = 1.0) -> AsyncGenerator[Dict[str, Any], None]:
        """
        Consume messages from specified topics.

        Args:
            topics: List of topics to consume from
            message_handler: Optional message handler function
            timeout: Poll timeout in seconds

        Yields:
            Message dictionaries
        """
        self.consumer = self.create_consumer()
        self.consumer.subscribe(topics)
        self.is_running = True

        while self.is_running:
            try:
                # Poll for messages
                msg = self.consumer.poll(timeout)

                if msg is None:
                    continue

                if msg.error():
                    if msg.error().code() == ConfluentKafkaError._PARTITION_EOF:
                        # End of partition
                        continue
                    else:
                        self.logger.error(f"Consumer error: {msg.error()}")
                        continue

                # Create message dictionary
                message = {
                    'topic': msg.topic(),
                    'partition': msg.partition(),
                    'offset': msg.offset(),
                    'key': msg.key().decode('utf-8') if msg.key() else None,
                    'value': msg.value().decode('utf-8') if msg.value() else None,
                    'timestamp': msg.timestamp(),
                    'headers': dict(msg.headers()) if msg.headers() else None,
                }

                # Call message handler if provided
                if message_handler:
                    message_handler(message)

                yield message

            except Exception as e:
                self.logger.error(f"Error consuming messages: {e}")
                raise ConsumerError(f"Error consuming messages: {e}", consumer_id=self.name)

    async def start_consumer(self, topics: List[str],
                           message_handler: Callable[[Dict[str, Any]], None] = None) -> asyncio.Task:
        """
        Start consuming messages in a background task.

        Args:
            topics: List of topics to consume from
            message_handler: Optional message handler function

        Returns:
            asyncio.Task for the consumer
        """
        # Clean up completed tasks
        self.tasks = [t for t in self.tasks if not t.done()]

        task = asyncio.create_task(
            self._consume_background(topics, message_handler)
        )
        self.tasks.append(task)
        return task

    async def _consume_background(self, topics: List[str],
                                message_handler: Callable[[Dict[str, Any]], None] = None):
        """Background consumer task."""
        try:
            async for message in self.consume_messages(topics, message_handler):
                if message_handler:
                    message_handler(message)
        except Exception as e:
            self.logger.error(f"Background consumer error: {e}")

    def stop(self):
        """Stop the consumer."""
        self.is_running = False

        # Cancel all tasks
        for task in self.tasks:
            if not task.done():
                task.cancel()

        self.tasks.clear()

        # Close consumer
        if self.consumer:
            self.consumer.close()

    def close(self):
        """Close the consumer."""
        self.stop()


class KafkaAdminClient:
    """Admin client for Kafka cluster operations."""

    def __init__(self, bootstrap_servers: str):
        self.admin_client = AdminClient({'bootstrap.servers': bootstrap_servers})
        self.logger = logging.getLogger(__name__)

    async def list_topics(self) -> Dict[str, Any]:
        """List all topics in the cluster."""
        try:
            metadata = self.admin_client.list_topics()
            topics_metadata = {}

            for topic, topic_metadata in metadata.topics.items():
                partitions_info = []
                if hasattr(topic_metadata, 'partitions'):
                    for partition_id, partition_info in topic_metadata.partitions.items():
                        partitions_info.append({
                            'id': partition_id,
                            'leader': partition_info.leader,
                            'replicas': partition_info.replicas,
                            'isr': partition_info.isrs
                        })

                topics_metadata[topic] = {
                    'partitions': [p['id'] for p in partitions_info],
                    'partitions_info': partitions_info
                }

            return {
                'topics': list(metadata.topics.keys()),
                'brokers': list(metadata.brokers),
                'topics_metadata': topics_metadata
            }
        except Exception as e:
            self.logger.error(f"Error listing topics: {e}")
            raise ConnectionError(f"Error listing topics: {e}")

    async def create_topic(self, topic_name: str, partitions: int = 3,
                          replication_factor: int = 1) -> bool:
        """Create a new topic."""
        try:
            new_topic = NewTopic(
                topic_name,
                num_partitions=partitions,
                replication_factor=replication_factor
            )

            result = self.admin_client.create_topics([new_topic])

            # Wait for operation to complete
            for topic, future in result.items():
                future.result()

            return True
        except Exception as e:
            self.logger.error(f"Error creating topic {topic_name}: {e}")
            raise ConnectionError(f"Error creating topic {topic_name}: {e}")

    async def describe_topic(self, topic_name: str) -> Dict[str, Any]:
        """Describe a specific topic."""
        try:
            metadata = self.admin_client.list_topics(topic_name)

            if topic_name not in metadata.topics:
                raise TopicError(f"Topic {topic_name} not found", topic_name=topic_name)

            topic_metadata = metadata.topics[topic_name]

            partitions_info = []
            if hasattr(topic_metadata, 'partitions'):
                for partition_id, partition_info in topic_metadata.partitions.items():
                    partitions_info.append({
                        'id': partition_id,
                        'leader': partition_info.leader,
                        'replicas': partition_info.replicas,
                        'isr': partition_info.isrs
                    })

            return {
                'topic': topic_name,
                'partitions': [p['id'] for p in partitions_info],
                'partitions_info': partitions_info
            }
        except Exception as e:
            self.logger.error(f"Error describing topic {topic_name}: {e}")
            raise TopicError(f"Error describing topic {topic_name}: {e}", topic_name=topic_name)

    async def delete_topic(self, topic_name: str) -> bool:
        """Delete a topic."""
        try:
            result = self.admin_client.delete_topics([topic_name])

            # Wait for operation to complete
            for topic, future in result.items():
                future.result()

            return True
        except Exception as e:
            self.logger.error(f"Error deleting topic {topic_name}: {e}")
            raise TopicError(f"Error deleting topic {topic_name}: {e}", topic_name=topic_name)

    async def list_consumer_groups(self) -> List[Dict[str, Any]]:
        """List all consumer groups."""
        try:
            future = self.admin_client.list_consumer_groups()
            # Wait for the future to complete
            result = await asyncio.get_event_loop().run_in_executor(None, lambda: future.result())
            groups = []

            for group in result.valid:
                groups.append({
                    'group_id': group.group_id,
                    'is_simple_consumer_group': group.is_simple_consumer_group,
                    'state': group.state,
                    'type': group.type
                })

            return groups
        except Exception as e:
            self.logger.error(f"Error listing consumer groups: {e}")
            raise ConsumerGroupError(f"Error listing consumer groups: {e}")

    async def describe_consumer_group(self, group_id: str) -> Dict[str, Any]:
        """Describe a specific consumer group."""
        try:
            # describe_consumer_groups returns a dict, not a Future
            metadata = self.admin_client.describe_consumer_groups([group_id])

            # Find the specific group in the valid list
            if group_id not in metadata.get('valid', {}):
                raise ConsumerGroupError(f"Group {group_id} not found", group_id=group_id)

            group = metadata['valid'][group_id]

            return {
                'group_id': group_id,
                'is_simple_consumer_group': group.is_simple_consumer_group,
                'state': group.state,
                'type': group.type
            }
        except Exception as e:
            self.logger.error(f"Error describing consumer group {group_id}: {e}")
            raise ConsumerGroupError(f"Error describing consumer group {group_id}: {e}", group_id=group_id)