"""
Kafka integration module for FastMCP.

This module provides the main integration point between Kafka functionality
and the FastMCP server.
"""

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

from mcp.server.fastmcp import FastMCP

from kafka_config import KafkaConfig, get_config, config_manager
from kafka_clients import AsyncProducer, ProducerPool, ConsumerManager, KafkaAdminClient, Message
from kafka_errors import ProducerError, ConsumerError, ConfigurationError, ConnectionError, ErrorHandler


@dataclass
class KafkaModuleConfig:
    """Configuration for KafkaModule."""
    bootstrap_servers: str
    max_pool_size: int = 5
    enable_producer_pool: bool = True
    enable_consumer_auto_start: bool = False


class KafkaModule:
    """Main Kafka integration module for FastMCP."""

    def __init__(self, mcp: FastMCP, config: Optional[KafkaConfig] = None):
        """Initialize Kafka module."""
        self.mcp = mcp
        self.config = config or get_config()
        self.logger = logging.getLogger(__name__)

        # Initialize components
        self.producer_pool = ProducerPool(max_pool_size=5) if self.config else None
        self.consumer_managers: Dict[str, ConsumerManager] = {}
        self.admin_client: Optional[KafkaAdminClient] = None

        # Track active producers
        self.active_producers: Dict[str, AsyncProducer] = {}

        # Initialize if config is available
        if self.config:
            self.initialize()

    def initialize(self):
        """Initialize the Kafka module with configuration."""
        if not self.config:
            raise ConfigurationError("Configuration not loaded")

        # Create admin client
        self.admin_client = KafkaAdminClient(
            bootstrap_servers=self.config.bootstrap_servers
        )

        self.logger.info(f"Kafka module initialized with bootstrap servers: {self.config.bootstrap_servers}")

    # Producer tools
    async def create_producer(self, config_name: str = 'default') -> AsyncProducer:
        """
        Create a producer instance.

        Args:
            config_name: Name of the producer configuration to use

        Returns:
            AsyncProducer instance
        """
        if not self.config:
            raise ConfigurationError("Configuration not loaded")

        # Get producer configuration from config manager
        producer_config = config_manager.get_producer_config(config_name)

        # Get producer from pool or create new one
        if self.producer_pool:
            kafka_producer = await self.producer_pool.get_producer(producer_config, config_name)
        else:
            # Create directly if pool is disabled
            from kafka_clients import Producer
            kafka_producer = Producer({
                'bootstrap.servers': producer_config.bootstrap_servers,
                'acks': producer_config.acks,
                'retries': producer_config.retries,
                'batch.size': producer_config.batch_size,
            })

        # Create async producer wrapper
        async_producer = AsyncProducer(kafka_producer, self.producer_pool, config_name)
        self.active_producers[config_name] = async_producer

        self.logger.info(f"Created producer with config '{config_name}'")
        return async_producer

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

        Args:
            topic: Topic name
            value: Message value (JSON string)
            key: Message key (optional)
            headers: Message headers (optional)
            producer_config: Name of producer configuration to use

        Returns:
            Dictionary with message delivery information
        """
        try:
            # Create producer if not exists
            if producer_config not in self.active_producers:
                producer = await self.create_producer(producer_config)
            else:
                producer = self.active_producers[producer_config]

            # Send message
            message = await producer.send_message(topic, value, key, headers)

            return {
                "success": True,
                "message": message.to_dict(),
                "producer_config": producer_config
            }

        except Exception as e:
            self.logger.error(f"Error sending message: {e}")
            return {
                "success": False,
                "error": ErrorHandler.handle_error(e),
                "details": ErrorHandler.get_error_details(e)
            }

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

        Args:
            messages: List of message dictionaries
            producer_config: Name of producer configuration to use

        Returns:
            Dictionary with batch delivery information
        """
        try:
            # Create producer if not exists
            if producer_config not in self.active_producers:
                producer = await self.create_producer(producer_config)
            else:
                producer = self.active_producers[producer_config]

            # Send messages batch
            results = await producer.send_messages_batch(messages)

            return {
                "success": True,
                "messages": [msg.to_dict() for msg in results],
                "count": len(results),
                "producer_config": producer_config
            }

        except Exception as e:
            self.logger.error(f"Error sending messages batch: {e}")
            return {
                "success": False,
                "error": ErrorHandler.handle_error(e),
                "details": ErrorHandler.get_error_details(e)
            }

    async def close_producer(self, producer_config: str = 'default') -> Dict[str, Any]:
        """
        Close a producer instance.

        Args:
            producer_config: Name of producer configuration to close

        Returns:
            Dictionary with closure status
        """
        try:
            if producer_config in self.active_producers:
                producer = self.active_producers[producer_config]
                await producer.close()
                del self.active_producers[producer_config]

                return {
                    "success": True,
                    "message": f"Producer '{producer_config}' closed successfully"
                }
            else:
                return {
                    "success": False,
                    "message": f"Producer '{producer_config}' not found"
                }

        except Exception as e:
            self.logger.error(f"Error closing producer: {e}")
            return {
                "success": False,
                "error": ErrorHandler.handle_error(e),
                "details": ErrorHandler.get_error_details(e)
            }

    # Consumer tools
    async def create_consumer(self, config_name: str = 'default') -> ConsumerManager:
        """
        Create a consumer instance.

        Args:
            config_name: Name of the consumer configuration to use

        Returns:
            ConsumerManager instance
        """
        if not self.config:
            raise ConfigurationError("Configuration not loaded")

        # Get consumer configuration from config manager
        consumer_config = config_manager.get_consumer_config(config_name)

        # Create consumer manager
        consumer_manager = ConsumerManager(consumer_config, config_name)
        self.consumer_managers[config_name] = consumer_manager

        self.logger.info(f"Created consumer with config '{config_name}'")
        return consumer_manager

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

        Args:
            topics: List of topics to consume from
            config_name: Name of consumer configuration to use
            message_handler: Optional message handler function
            timeout: Poll timeout in seconds

        Returns:
            Dictionary with consumer status
        """
        try:
            # Create consumer if not exists
            if config_name not in self.consumer_managers:
                consumer = await self.create_consumer(config_name)
            else:
                consumer = self.consumer_managers[config_name]

            # Start consuming messages
            consumer_task = asyncio.create_task(
                consumer.start_consumer(topics, message_handler)
            )

            return {
                "success": True,
                "message": f"Started consuming messages from topics: {topics}",
                "consumer_config": config_name,
                "task_id": id(consumer_task)
            }

        except Exception as e:
            self.logger.error(f"Error consuming messages: {e}")
            return {
                "success": False,
                "error": ErrorHandler.handle_error(e),
                "details": ErrorHandler.get_error_details(e)
            }

    async def close_consumer(self, config_name: str = 'default') -> Dict[str, Any]:
        """
        Close a consumer instance.

        Args:
            config_name: Name of consumer configuration to close

        Returns:
            Dictionary with closure status
        """
        try:
            if config_name in self.consumer_managers:
                consumer = self.consumer_managers[config_name]
                consumer.close()
                del self.consumer_managers[config_name]

                return {
                    "success": True,
                    "message": f"Consumer '{config_name}' closed successfully"
                }
            else:
                return {
                    "success": False,
                    "message": f"Consumer '{config_name}' not found"
                }

        except Exception as e:
            self.logger.error(f"Error closing consumer: {e}")
            return {
                "success": False,
                "error": ErrorHandler.handle_error(e),
                "details": ErrorHandler.get_error_details(e)
            }

    # Topic management tools
    async def list_topics(self) -> Dict[str, Any]:
        """
        List all topics in the cluster.

        Returns:
            Dictionary with topic information
        """
        try:
            if not self.admin_client:
                raise ConfigurationError("Admin client not initialized")

            topics_data = await self.admin_client.list_topics()

            return {
                "success": True,
                "topics": topics_data
            }

        except Exception as e:
            self.logger.error(f"Error listing topics: {e}")
            return {
                "success": False,
                "error": ErrorHandler.handle_error(e),
                "details": ErrorHandler.get_error_details(e)
            }

    async def create_topic(self, topic_name: str, partitions: int = 3,
                          replication_factor: int = 1) -> Dict[str, Any]:
        """
        Create a new topic.

        Args:
            topic_name: Name of the topic to create
            partitions: Number of partitions
            replication_factor: Replication factor

        Returns:
            Dictionary with creation status
        """
        try:
            if not self.admin_client:
                raise ConfigurationError("Admin client not initialized")

            success = await self.admin_client.create_topic(topic_name, partitions, replication_factor)

            return {
                "success": True,
                "message": f"Topic '{topic_name}' created successfully",
                "topic": topic_name,
                "partitions": partitions,
                "replication_factor": replication_factor
            }

        except Exception as e:
            self.logger.error(f"Error creating topic: {e}")
            return {
                "success": False,
                "error": ErrorHandler.handle_error(e),
                "details": ErrorHandler.get_error_details(e)
            }

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

        Args:
            topic_name: Name of the topic to describe

        Returns:
            Dictionary with topic details
        """
        try:
            if not self.admin_client:
                raise ConfigurationError("Admin client not initialized")

            topic_info = await self.admin_client.describe_topic(topic_name)

            return {
                "success": True,
                "topic": topic_info
            }

        except Exception as e:
            self.logger.error(f"Error describing topic: {e}")
            return {
                "success": False,
                "error": ErrorHandler.handle_error(e),
                "details": ErrorHandler.get_error_details(e)
            }

    async def delete_topic(self, topic_name: str) -> Dict[str, Any]:
        """
        Delete a topic.

        Args:
            topic_name: Name of the topic to delete

        Returns:
            Dictionary with deletion status
        """
        try:
            if not self.admin_client:
                raise ConfigurationError("Admin client not initialized")

            success = await self.admin_client.delete_topic(topic_name)

            return {
                "success": True,
                "message": f"Topic '{topic_name}' deleted successfully"
            }

        except Exception as e:
            self.logger.error(f"Error deleting topic: {e}")
            return {
                "success": False,
                "error": ErrorHandler.handle_error(e),
                "details": ErrorHandler.get_error_details(e)
            }

    # Consumer group management tools
    async def list_consumer_groups(self) -> Dict[str, Any]:
        """
        List all consumer groups.

        Returns:
            Dictionary with consumer group information
        """
        try:
            if not self.admin_client:
                raise ConfigurationError("Admin client not initialized")

            groups = await self.admin_client.list_consumer_groups()

            return {
                "success": True,
                "consumer_groups": groups
            }

        except Exception as e:
            self.logger.error(f"Error listing consumer groups: {e}")
            return {
                "success": False,
                "error": ErrorHandler.handle_error(e),
                "details": ErrorHandler.get_error_details(e)
            }

    async def describe_consumer_group(self, group_id: str) -> Dict[str, Any]:
        """
        Describe a specific consumer group.

        Args:
            group_id: ID of the consumer group to describe

        Returns:
            Dictionary with consumer group details
        """
        try:
            if not self.admin_client:
                raise ConfigurationError("Admin client not initialized")

            group_info = await self.admin_client.describe_consumer_group(group_id)

            return {
                "success": True,
                "consumer_group": group_info
            }

        except Exception as e:
            self.logger.error(f"Error describing consumer group: {e}")
            return {
                "success": False,
                "error": ErrorHandler.handle_error(e),
                "details": ErrorHandler.get_error_details(e)
            }

    # System status resource
    def get_status(self) -> Dict[str, Any]:
        """
        Get module status information.

        Returns:
            Dictionary with status information
        """
        status = {
            "success": True,
            "module": "kafka",
            "status": "active",
            "config_loaded": self.config is not None,
            "bootstrap_servers": self.config.bootstrap_servers if self.config else None,
            "active_producers": len(self.active_producers),
            "active_consumers": len(self.consumer_managers),
            "producer_pool_active": self.producer_pool is not None
        }

        if self.producer_pool:
            status["producer_pool_size"] = sum(
                len(pool) for pool in self.producer_pool.producers.values()
            )

        return status

    def close_all(self):
        """Close all Kafka connections."""
        self.logger.info("Closing all Kafka connections")

        # Close all producers
        for producer in self.active_producers.values():
            producer.close()
        self.active_producers.clear()

        # Close all consumers
        for consumer in self.consumer_managers.values():
            consumer.close()
        self.consumer_managers.clear()

        # Close producer pool
        if self.producer_pool:
            self.producer_pool.close_all()

        self.logger.info("All Kafka connections closed")


def create_kafka_module(mcp: FastMCP, config: Optional[KafkaConfig] = None) -> KafkaModule:
    """
    Create and initialize a Kafka module.

    Args:
        mcp: FastMCP instance
        config: Optional Kafka configuration

    Returns:
        KafkaModule instance
    """
    kafka_module = KafkaModule(mcp, config)
    return kafka_module