"""
Kafka tools module for FastMCP.

This module defines all the MCP tools for Kafka operations.
"""

import json
import logging
from typing import Dict, List, Any, Optional, Callable

from mcp.server.fastmcp import Context

from kafka_integration import KafkaModule


def register_kafka_tools(kafka_module: KafkaModule, mcp: Any):
    """
    Register all Kafka tools with FastMCP.

    Args:
        kafka_module: KafkaModule instance
        mcp: FastMCP instance
    """

    @mcp.tool()
    async def create_producer(ctx: Context, config_name: str = 'default') -> Dict[str, Any]:
        """
        Create a Kafka producer instance for sending messages.

        Args:
            config_name: Name of the producer configuration to use (default: 'default')

        Returns:
            Dictionary with producer creation status and configuration
        """
        try:
            producer = await kafka_module.create_producer(config_name)
            return {
                "success": True,
                "message": f"Producer '{config_name}' created successfully",
                "config_name": config_name,
                "producer_id": config_name
            }
        except Exception as e:
            return {
                "success": False,
                "error": str(e),
                "config_name": config_name
            }

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

        Args:
            topic: Kafka topic name
            value: Message value (JSON string or plain text)
            key: Optional message key for partitioning
            headers: Optional JSON string of message headers
            producer_config: Name of producer configuration to use (default: 'default')

        Returns:
            Dictionary with message delivery confirmation
        """
        try:
            # Parse headers if provided
            header_dict = None
            if headers:
                try:
                    header_dict = json.loads(headers)
                    # Convert string values to bytes for Kafka headers
                    header_dict = {k: str(v) for k, v in header_dict.items()}
                except json.JSONDecodeError:
                    return {
                        "success": False,
                        "error": "Invalid JSON format for headers",
                        "topic": topic
                    }

            result = await kafka_module.send_message(topic, value, key, header_dict, producer_config)
            return result

        except Exception as e:
            return {
                "success": False,
                "error": str(e),
                "topic": topic
            }

    @mcp.tool()
    async def send_messages_batch(
        ctx: Context,
        messages_json: str,
        producer_config: str = 'default'
    ) -> Dict[str, Any]:
        """
        Send multiple messages to Kafka in a batch.

        Args:
            messages_json: JSON string containing array of message objects with keys:
                - topic (required): Kafka topic name
                - value (required): Message value
                - key (optional): Message key for partitioning
                - headers (optional): Message headers as dictionary
            producer_config: Name of producer configuration to use (default: 'default')

        Returns:
            Dictionary with batch delivery results

        Example messages_json format:
            [
                {"topic": "my-topic", "value": "Hello", "key": "key1"},
                {"topic": "my-topic", "value": "World", "key": "key2"}
            ]
        """
        try:
            # Parse messages JSON
            messages = json.loads(messages_json)

            if not isinstance(messages, list):
                return {
                    "success": False,
                    "error": "Messages must be a JSON array"
                }

            # Validate each message
            for i, msg in enumerate(messages):
                if not isinstance(msg, dict):
                    return {
                        "success": False,
                        "error": f"Message at index {i} is not an object"
                    }
                if 'topic' not in msg or 'value' not in msg:
                    return {
                        "success": False,
                        "error": f"Message at index {i} missing required 'topic' or 'value' field"
                    }

            result = await kafka_module.send_messages_batch(messages, producer_config)
            return result

        except json.JSONDecodeError:
            return {
                "success": False,
                "error": "Invalid JSON format for messages"
            }
        except Exception as e:
            return {
                "success": False,
                "error": str(e)
            }

    @mcp.tool()
    async def close_producer(ctx: Context, producer_config: str = 'default') -> Dict[str, Any]:
        """
        Close a Kafka producer instance and release resources.

        Args:
            producer_config: Name of producer configuration to close (default: 'default')

        Returns:
            Dictionary with closure status
        """
        result = await kafka_module.close_producer(producer_config)
        return result

    @mcp.tool()
    async def create_consumer(ctx: Context, config_name: str = 'default') -> Dict[str, Any]:
        """
        Create a Kafka consumer instance for reading messages.

        Args:
            config_name: Name of the consumer configuration to use (default: 'default')

        Returns:
            Dictionary with consumer creation status
        """
        try:
            consumer = await kafka_module.create_consumer(config_name)
            return {
                "success": True,
                "message": f"Consumer '{config_name}' created successfully",
                "config_name": config_name,
                "consumer_id": config_name
            }
        except Exception as e:
            return {
                "success": False,
                "error": str(e),
                "config_name": config_name
            }

    @mcp.tool()
    async def consume_messages(
        ctx: Context,
        topics: List[str],
        config_name: str = 'default',
        timeout: float = 1.0
    ) -> Dict[str, Any]:
        """
        Start consuming messages from specified Kafka topics.

        Note: This starts a background consumer task. Messages will be consumed asynchronously.
        Use close_consumer to stop the consumer.

        Args:
            topics: List of Kafka topic names to consume from
            config_name: Name of consumer configuration to use (default: 'default')
            timeout: Poll timeout in seconds (default: 1.0)

        Returns:
            Dictionary with consumer start status and task ID
        """
        try:
            if not isinstance(topics, list):
                return {
                    "success": False,
                    "error": "Topics must be a list of strings"
                }

            if len(topics) == 0:
                return {
                    "success": False,
                    "error": "At least one topic must be specified"
                }

            result = await kafka_module.consume_messages(topics, config_name, None, timeout)
            return result

        except Exception as e:
            return {
                "success": False,
                "error": str(e),
                "topics": topics
            }

    @mcp.tool()
    async def close_consumer(ctx: Context, config_name: str = 'default') -> Dict[str, Any]:
        """
        Close a Kafka consumer instance and stop consuming messages.

        Args:
            config_name: Name of consumer configuration to close (default: 'default')

        Returns:
            Dictionary with closure status
        """
        result = await kafka_module.close_consumer(config_name)
        return result

    @mcp.tool()
    async def list_topics(ctx: Context) -> Dict[str, Any]:
        """
        List all available topics in the Kafka cluster.

        Returns:
            Dictionary with topic list and metadata
        """
        result = await kafka_module.list_topics()
        return result

    @mcp.tool()
    async def create_topic(
        ctx: Context,
        topic_name: str,
        partitions: int = 3,
        replication_factor: int = 1
    ) -> Dict[str, Any]:
        """
        Create a new Kafka topic with specified configuration.

        Args:
            topic_name: Name for the new topic
            partitions: Number of partitions for the topic (default: 3)
            replication_factor: Replication factor for the topic (default: 1)

        Returns:
            Dictionary with topic creation status
        """
        try:
            if partitions < 1:
                return {
                    "success": False,
                    "error": "Partitions must be at least 1"
                }

            if replication_factor < 1:
                return {
                    "success": False,
                    "error": "Replication factor must be at least 1"
                }

            result = await kafka_module.create_topic(topic_name, partitions, replication_factor)
            return result

        except Exception as e:
            return {
                "success": False,
                "error": str(e),
                "topic_name": topic_name
            }

    @mcp.tool()
    async def describe_topic(ctx: Context, topic_name: str) -> Dict[str, Any]:
        """
        Get detailed information about a specific Kafka topic.

        Args:
            topic_name: Name of the topic to describe

        Returns:
            Dictionary with topic details including partitions, replicas, and ISR
        """
        result = await kafka_module.describe_topic(topic_name)
        return result

    @mcp.tool()
    async def delete_topic(ctx: Context, topic_name: str) -> Dict[str, Any]:
        """
        Delete a Kafka topic from the cluster.

        Warning: This operation is irreversible and will delete all messages in the topic.

        Args:
            topic_name: Name of the topic to delete

        Returns:
            Dictionary with topic deletion status
        """
        result = await kafka_module.delete_topic(topic_name)
        return result

    @mcp.tool()
    async def list_consumer_groups(ctx: Context) -> Dict[str, Any]:
        """
        List all consumer groups in the Kafka cluster.

        Returns:
            Dictionary with consumer group list and metadata
        """
        result = await kafka_module.list_consumer_groups()
        return result

    @mcp.tool()
    async def describe_consumer_group(ctx: Context, group_id: str) -> Dict[str, Any]:
        """
        Get detailed information about a specific consumer group.

        Args:
            group_id: ID of the consumer group to describe

        Returns:
            Dictionary with consumer group details including members and assignments
        """
        result = await kafka_module.describe_consumer_group(group_id)
        return result

    logging.info("Kafka tools registered successfully")
