KafkaV3

from collections.abc import Callable
from io import StringIO

import demistomock as demisto
from CommonServerPython import *
from confluent_kafka import TIMESTAMP_NOT_AVAILABLE, Consumer, KafkaException, Message, Producer, TopicPartition
from confluent_kafka.schema_registry import SchemaRegistryClient
from confluent_kafka.schema_registry.avro import AvroSerializer
from confluent_kafka.serialization import MessageField, SerializationContext

""" IMPORTS """
import json
import logging
import tempfile
import traceback

import urllib3

# Disable insecure warnings
urllib3.disable_warnings()

SUPPORTED_GENERAL_OFFSETS = ["smallest", "earliest", "beginning", "largest", "latest", "end", "error"]

""" CLIENT CLASS """


class KConsumer(Consumer):
    """Empty inheritance class for C-typed class in order to make mocking work."""


class KProducer(Producer):
    """Empty inheritance class for C-typed class in order to make mocking work."""


class KSchemaRegistryClient(SchemaRegistryClient):
    """Empty inheritance class for C-typed class in order to make mocking work."""


class KafkaCommunicator:
    """Client class to interact with Kafka."""

    conf_producer: Optional[dict[str, Any]] = None
    conf_consumer: Optional[dict[str, Any]] = None
    conf_schema_registry: Optional[dict[str, Any]] = None
    ca_path: Optional[str] = None
    client_cert_path: Optional[str] = None
    client_key_path: Optional[str] = None
    kafka_logger: Optional[logging.Logger] = None

    SESSION_TIMEOUT: int = 10000
    REQUESTS_TIMEOUT: float = 10.0
    POLL_TIMEOUT: float = 1.0
    POLL_TIMEOUT_STOP_UPON_TIMEOUT = 10.0
    # which caused test playbook failures in builds.
    MAX_POLLS_FOR_LOG: int = 100

    def __init__(
        self,
        brokers: str,
        use_ssl: bool,
        plain_password: Optional[str] = None,
        plain_username: Optional[str] = None,
        use_sasl: bool = False,
        offset: str = "earliest",
        group_id: str = "xsoar_group",
        message_max_bytes: Optional[int] = None,
        ca_cert: Optional[str] = None,
        client_cert: Optional[str] = None,
        client_cert_key: Optional[str] = None,
        ssl_password: Optional[str] = None,
        trust_any_cert: bool = False,
        kafka_logger: Optional[logging.Logger] = None,
        schema_registry_url: Optional[str] = None,
        schema_registry_username: Optional[str] = None,
        schema_registry_password: Optional[str] = None,
    ):
        """Set configuration dicts for consumer and producer.

        Args:
            brokers: The broker ips:ports to connect to.
            offset: The offset to start consuming from.
            group_id: The consumer group id.
            message_max_bytes: The maximum bytes a message can have.
            ca_cert: The contents of the CA certificate.
            client_cert: The contents of the client certificate.
            client_cert_key: The contents of the client certificate's key
            ssl_password: The password with which the client certificate is protected by.
            schema_registry_url: The URL of the schema registry.
            schema_registry_username: The username for the schema registry.
            schema_registry_password: The password for the schema registry.
        """

        # Set producer conf dict
        self.conf_producer = {}
        self.update_client_dict(
            self.conf_producer,
            trust_any_cert,
            use_ssl,
            ca_cert,
            client_cert,
            client_cert_key,
            ssl_password,
            use_sasl,
            plain_username,
            plain_password,
            brokers,
        )

        if offset not in SUPPORTED_GENERAL_OFFSETS:
            raise DemistoException(f"General offset {offset} not found in supported offsets: {SUPPORTED_GENERAL_OFFSETS}")

        # Set consumer conf dict
        self.conf_consumer = {
            "session.timeout.ms": self.SESSION_TIMEOUT,
            "auto.offset.reset": offset,
            "group.id": group_id,
            "enable.auto.commit": False,
        }

        self.update_client_dict(
            self.conf_consumer,
            trust_any_cert,
            use_ssl,
            ca_cert,
            client_cert,
            client_cert_key,
            ssl_password,
            use_sasl,
            plain_username,
            plain_password,
            brokers,
        )

        self.kafka_logger = kafka_logger

        if message_max_bytes:
            self.conf_consumer.update({"message.max.bytes": int(message_max_bytes)})

        demisto.debug(f"The consumer configuration is \n{self.conf_consumer}\n")
        demisto.debug(f"The producer configuration is \n{self.conf_producer}\n")

        # Set schema registry conf dict
        if schema_registry_url:
            self.conf_schema_registry = {
                "url": schema_registry_url,
            }
            if schema_registry_username and schema_registry_password:
                self.conf_schema_registry["basic.auth.user.info"] = f"{schema_registry_username}:{schema_registry_password}"

            demisto.debug(f"The schema registry configuration is  \n{self.conf_schema_registry}\n")

    def update_client_dict(
        self,
        client_dict,
        trust_any_cert,
        use_ssl,
        ca_cert,
        client_cert,
        client_cert_key,
        ssl_password,
        use_sasl,
        plain_username,
        plain_password,
        brokers,
    ):
        """
        Updates the conf_producer or conf_consumer configuration based on the specified authentication method.
        It assumes that all required parameters have been validated by the validate_params function.

        Args:
            client_dict (dict): The configuration dictionary to be updated.
            This should be either the producer's or consumer's conf_dict.
        """

        client_dict.update({"bootstrap.servers": brokers})

        if use_ssl and not use_sasl:
            if self.ca_path:
                client_dict.update(
                    {
                        "ssl.ca.location": self.ca_path,
                        "ssl.certificate.location": self.client_cert_path,
                        "ssl.key.location": self.client_key_path,
                        "security.protocol": "ssl",
                    }
                )
            else:
                # temporary creating ca certification file
                with tempfile.NamedTemporaryFile(mode="w", delete=False) as ca_descriptor:
                    self.ca_path = ca_descriptor.name
                    ca_descriptor.write(ca_cert)
                client_dict.update({"ssl.ca.location": self.ca_path})

                # temporary creating client certification file
                with tempfile.NamedTemporaryFile(mode="w", delete=False) as client_cert_descriptor:
                    self.client_cert_path = client_cert_descriptor.name
                    client_cert_descriptor.write(client_cert)
                client_dict.update({"ssl.certificate.location": self.client_cert_path})

                # temporary creating client certification's key file
                with tempfile.NamedTemporaryFile(mode="w", delete=False) as client_key_descriptor:
                    self.client_key_path = client_key_descriptor.name
                    client_key_descriptor.write(client_cert_key)
                client_dict.update({"ssl.key.location": self.client_key_path, "security.protocol": "ssl"})

            if ssl_password:
                client_dict.update({"ssl.key.password": ssl_password})

        # SASL with SSL
        elif use_sasl:
            client_dict.update(
                {
                    "security.protocol": "SASL_SSL",
                    "sasl.mechanism": "PLAIN",
                    "sasl.username": plain_username,
                    "sasl.password": plain_password,
                }
            )

            if trust_any_cert:
                client_dict.update(
                    {"ssl.endpoint.identification.algorithm": "none", "enable.ssl.certificate.verification": False}
                )
            else:  # ca_cert
                if self.ca_path:
                    client_dict.update({"ssl.ca.location": self.ca_path})
                else:
                    with tempfile.NamedTemporaryFile(mode="w", delete=False) as ca_descriptor:
                        self.ca_path = ca_descriptor.name
                        ca_descriptor.write(ca_cert)
                    client_dict.update({"ssl.ca.location": self.ca_path})

                if ssl_password:
                    client_dict.update({"ssl.key.password": ssl_password})

    def get_kafka_consumer(self) -> KConsumer:
        if self.kafka_logger:
            return KConsumer(self.conf_consumer, logger=self.kafka_logger)
        return KConsumer(self.conf_consumer)

    def get_kafka_producer(self) -> KProducer:
        if self.kafka_logger:
            return KProducer(self.conf_producer, logger=self.kafka_logger)
        return KProducer(self.conf_producer)

    def get_kafka_schema_registry(self) -> Optional[KSchemaRegistryClient]:
        if self.conf_schema_registry:
            return KSchemaRegistryClient(self.conf_schema_registry)
        return None

    def update_conf_for_fetch(self, message_max_bytes: Optional[int] = None):
        """Update consumer configurations for fetching messages

        Args:
            message_max_bytes: The maximum message bytes to fetch

        Raise DemistoException if consumer was not initialized before.
        """
        if self.conf_consumer:
            if message_max_bytes:
                self.conf_consumer.update({"message.max.bytes": int(message_max_bytes)})

        else:
            raise DemistoException("Kafka consumer was not yet initialized.")

    def test_connection(self, log_stream: Optional[StringIO] = None) -> str:
        """Test getting topics with the consumer and producer configurations."""
        error_msg = ""
        consumer: Optional[KConsumer] = None
        producer: Optional[KProducer] = None
        schema_registry: Optional[KSchemaRegistryClient] = None

        try:
            consumer = self.get_kafka_consumer()
            consumer_topics = consumer.list_topics(timeout=self.REQUESTS_TIMEOUT)
            consumer_topics.topics  # noqa: B018

        except Exception as e:
            error_msg = f"Error connecting to kafka using consumer: {e!s}\n{traceback.format_exc()}"
        finally:
            try:
                if error_msg and consumer:
                    demisto.debug("Polling consumer for debug logs")
                    consumer.poll(1.0)  # For the logger to be updated with consumer errors.
                    if log_stream:
                        polls = 0
                        while not log_stream.getvalue() and polls < self.MAX_POLLS_FOR_LOG:
                            polls += 1
                            consumer.poll(1.0)
            finally:
                if error_msg:
                    raise DemistoException(error_msg)

        try:
            producer = self.get_kafka_producer()
            producer_topics = producer.list_topics(timeout=self.REQUESTS_TIMEOUT)
            producer_topics.topics  # noqa: B018

        except Exception as e:
            error_msg = f"Error connecting to kafka using producer: {e!s}\n{traceback.format_exc()}"
        finally:
            try:
                if error_msg and producer:
                    demisto.debug("Polling producer for debug logs")
                    producer.flush()  # For the logger to updated with producer errors.
                    if log_stream:
                        polls = 0
                        while not log_stream.getvalue() and polls < self.MAX_POLLS_FOR_LOG:
                            polls += 1
                            producer.flush()
            finally:
                if error_msg:
                    raise DemistoException(error_msg)

        try:
            schema_registry = self.get_kafka_schema_registry()
            if schema_registry:
                schema_registry.get_subjects()

        except Exception as e:
            raise DemistoException(f"Error connecting to kafka schema registry: {e!s}\n{traceback.format_exc()}")

        return "ok"

    @staticmethod
    def delivery_report(err: KafkaException, msg: Message) -> None:
        """Callback function for producer. It is called when producer.flush() is called."""
        if err is not None:
            demisto.debug(f"Kafka v3 - Message {msg} delivery failed: {err}")
            raise DemistoException(f"Message delivery failed: {err}")
        else:
            return_results(f"Message was successfully produced to topic '{msg.topic()}', partition {msg.partition()}")

    def get_topics(self, consumer: bool = True) -> dict:
        """Get Kafka topics

        Args:
            consumer: Whether to get topics using the consumer configuration (True) or producer configuration (False).

        Return topics metadata object as described in confluent-kafka API
        """
        if consumer:
            client = self.get_kafka_consumer()
        else:
            client = self.get_kafka_producer()
        cluster_metadata = client.list_topics(timeout=self.REQUESTS_TIMEOUT)
        return cluster_metadata.topics

    def get_partition_offsets(self, topic: str, partition: int) -> tuple[int, int]:
        """Get earliest and latest offsets for the specified partition in the specified topic.

        Return (earliest offset, latest offset)
        """
        kafka_consumer = self.get_kafka_consumer()
        partition = TopicPartition(topic=topic, partition=partition)
        return kafka_consumer.get_watermark_offsets(partition=partition, timeout=self.REQUESTS_TIMEOUT)

    def produce(
        self,
        topic: str,
        value: str,
        value_schema_type: Optional[str],
        value_schema_str: Optional[str],
        value_schema_subject_name: Optional[str],
        partition: int,
    ) -> None:
        """Produce in to kafka

        Args:
            topic: The topic to produce to
            value: The message/object to write
            value_schema_type: The schema type of the value
            value_schema_str: The schema str of the value
            value_schema_subject_name: The schema subject name of the value
            partition: The partition to produce to.

        The delivery_report is called after production.
        """
        kafka_producer = self.get_kafka_producer()
        serialized_value = value

        if value_schema_type:
            kafka_schema_registry_client = self.get_kafka_schema_registry()
            if not kafka_schema_registry_client:
                raise DemistoException(
                    "Kafka Schema Registry client is not configured. Please configure one to use schema validation."
                )
            if not value_schema_str and not value_schema_subject_name:
                raise DemistoException("Schema is not provided. Please provide one.")
            if value_schema_str and value_schema_subject_name:
                raise DemistoException(
                    "Both value_schema_str and value_schema_subject_name are provided. Please provide only one."
                )

            resolved_schema_str = value_schema_str
            # Retrieve schema from schema registry
            if value_schema_subject_name:
                registered_schema = kafka_schema_registry_client.get_latest_version(subject_name=value_schema_subject_name)
                if registered_schema.schema.schema_type != value_schema_type:
                    raise DemistoException(
                        f"Unsupported schema type '{registered_schema.schema.schema_type}'. Expected '{value_schema_type}'."
                    )
                resolved_schema_str = registered_schema.schema.schema_str

            if value_schema_type == "AVRO":
                avro_serializer = AvroSerializer(
                    schema_str=resolved_schema_str, schema_registry_client=kafka_schema_registry_client
                )
                serialized_value = avro_serializer(json.loads(value), SerializationContext(topic, MessageField.VALUE))

        kafka_producer.produce(topic=topic, value=serialized_value, partition=partition, on_delivery=self.delivery_report)
        kafka_producer.flush()

    def consume(self, poll_timeout: float, topic: str, partition: int = -1, offset: str = "0") -> Message:
        """Consume a message from kafka

        Args:
            topic: The topic to consume from
            partition: The partition to consume from
            offset: The offset to start consuming from

        Return the consumed kafka message in the confluent_kafka API's format.
        """
        kafka_consumer = self.get_kafka_consumer()
        kafka_consumer.assign(self.get_topic_partitions(topic, partition, offset, True))
        polled_msg = kafka_consumer.poll(poll_timeout)
        demisto.debug(f"polled {polled_msg} with {poll_timeout=}")
        kafka_consumer.close()
        return polled_msg

    def get_offset_for_partition(self, topic: str, partition: int, offset: int | str) -> int:
        """Get the numerical offset from a partition of a topic

        Args:
            topic: The relavant topic
            partition: The relevant partition
            offset: Which offset to retrieve 'earliest'/ 'latest'

        Return the numerical value of the specified offset.
        """
        earliest_offset, oldest_offset = self.get_partition_offsets(topic=topic, partition=partition)
        offset = str(offset)
        if offset.lower() == "earliest":
            return earliest_offset
        elif offset.lower() == "latest":
            return oldest_offset - 1
        else:
            number_offset = int(offset)
            if number_offset < int(earliest_offset) or number_offset >= int(oldest_offset):
                raise DemistoException(
                    f"Offset {offset} for topic {topic} and partition {partition} is out of bounds "
                    f"[{earliest_offset}, {oldest_offset})"
                )
            return number_offset

    @logger
    def get_topic_partitions(self, topic: str, partition: int | list, offset: str | int, consumer: bool = False) -> list:
        """Get relevant TopicPartiton structures to specify for the consumer.

        Args:
            topic: The topic
            partition: The partition
            offset: The offset could be either int or 'earliest' / 'latest'
            consumer: True to use consumer configuration when connecting to kafka False for producer configuration.

        Return a list of TopicPartition objects, ready for consumer assign command.
        """
        topic_partitions = []
        if partition != -1 and not isinstance(partition, list):
            demisto.debug(f"Got single partition {partition}, getting offsets with offset {offset}")
            updated_offset = self.get_offset_for_partition(topic, int(partition), offset)
            topic_partitions = [TopicPartition(topic=topic, partition=int(partition), offset=updated_offset)]

        elif isinstance(partition, list):
            demisto.debug(f"Got partition list {partition}, getting offsets with offset {offset}")
            for single_partition in partition:
                try:
                    updated_offset = self.get_offset_for_partition(topic, int(single_partition), offset)
                    topic_partitions += [TopicPartition(topic=topic, partition=int(single_partition), offset=updated_offset)]
                except KafkaException as e:
                    # Sometimes listing topics can return uninitialized partitions.
                    # If that's the case, ignore them and continue.
                    if "Unknown partition" not in str(e):
                        raise e

        else:
            topics = self.get_topics(consumer=consumer)
            topic_metadata = topics[topic]
            demisto.debug(f"Got no partition, getting all partitions and offsets with offset {offset}")
            for metadata_partition in topic_metadata.partitions.values():
                try:
                    updated_offset = self.get_offset_for_partition(topic, metadata_partition.id, offset)
                    topic_partitions += [TopicPartition(topic=topic, partition=metadata_partition.id, offset=updated_offset)]
                except KafkaException as e:
                    # Sometimes listing topics can return uninitialized partitions.
                    # If that's the case, ignore them and continue.
                    if "Unknown partition" not in str(e):
                        raise e

        return topic_partitions


""" HELPER FUNCTIONS """


def validate_params(
    use_ssl, use_sasl, trust_any_cert, plain_username, plain_password, brokers, ca_cert, client_cert, client_cert_key
):
    """
    The function validates parameters for SSL and SASL_SSL authentication methods and raises an error if any invalid
    configurations are detected.

    For SSL authentication, it checks if use_ssl is True and requires ca_cert, client_cert, and client_cert_key parameters.

    For SASL_SSL authentication, it checks if use_sasl is True and requires plain_username, plain_password

    The brokers parameter is mandatory for both authentication methods.
    """

    # Check if brokers are provided
    if not brokers:
        raise DemistoException("Please specify a CSV list of Kafka brokers to connect to.")

    # Helper function to check for missing parameters

    def check_missing_params(params, missing):
        for param, param_name in params:
            if not param:
                missing.append(param_name)

    missing: List[str] = []

    # Check SSL requirements
    if use_ssl:
        ssl_params = [
            (ca_cert, "CA certificate of Kafka server (.cer)"),
            (client_cert, "Client certificate (.cer)"),
            (client_cert_key, "Client certificate key (.key)"),
        ]
        check_missing_params(ssl_params, missing)

    # Check SASL_PLAIN requirements
    if use_sasl:
        sasl_params = [(plain_username, "SASL PLAIN Username"), (plain_password, "SASL PLAIN Password")]
        if not trust_any_cert:
            sasl_params.append((ca_cert, "CA certificate of Kafka server (.cer)"))

        check_missing_params(sasl_params, missing)

    if missing:
        missing_items = ", ".join(missing)
        raise DemistoException(f"Missing required parameters: {missing_items}. Please provide them.")


def capture_logs(func: Callable):
    """Capture confluent kafka logs and add them when raising exceptions.

    Args:
        func: Has to support kafka_logger and log_stream kwargs

    return: the func's result
    """

    def wrapper(*args, **kwargs):
        logging.raiseExceptions = False
        log_stream = StringIO()
        log_handler = logging.StreamHandler(stream=log_stream)
        log_handler.setFormatter(logging.Formatter(logging.BASIC_FORMAT))
        kafka_logger = logging.getLogger()
        kafka_logger.addHandler(log_handler)
        kafka_logger.setLevel(logging.DEBUG)
        debug_log = ""

        try:
            kwargs["kafka_logger"] = kafka_logger
            kwargs["log_stream"] = log_stream
            result = func(*args, **kwargs)
        except Exception as e:
            captured_logs = log_stream.getvalue()
            if captured_logs:
                debug_log = f"\n{e!s}\nDebug logs:\n{captured_logs}\n"
            else:
                debug_log = f"\n{e!s}\n"
        finally:
            if log_stream:
                try:
                    logging.getLogger().removeHandler(log_handler)
                    log_stream.close()
                except Exception as e:
                    debug_log = (
                        f"Kafka v3: unexpected exception when trying to remove log handler:{e}\n\n"
                        f"Other Exceptions:{debug_log}"
                    )
                finally:
                    if debug_log:
                        raise DemistoException(debug_log)

        return result

    return wrapper


def create_incident(message: Message, topic: str) -> dict:
    """Create incident from kafka's message.

    Args:
        message: Kafka message to create incident from
        topic: Message's topic

    Return incident
    """
    message_value = message.value()
    raw = {"Topic": topic, "Partition": message.partition(), "Offset": message.offset(), "Message": message_value.decode("utf-8")}
    incident = {
        "name": f"Kafka {topic} partition:{message.partition()} offset:{message.offset()}",
        "details": message_value.decode("utf-8"),
        "rawJSON": json.dumps(raw),
    }

    timestamp = message.timestamp()  # returns a list of [timestamp_type, timestamp]
    if timestamp and len(timestamp) == 2 and timestamp[0] != TIMESTAMP_NOT_AVAILABLE:
        incident["occurred"] = timestamp_to_datestring(timestamp[1])

    demisto.debug(f"Creating incident from topic {topic} partition {message.partition()} offset {message.offset()}")
    return incident


""" COMMANDS """


def command_test_module(kafka: KafkaCommunicator, demisto_params: dict, log_stream: Optional[StringIO] = None) -> str:
    """Test getting available topics using consumer and producer configurations.
    Validate the fetch parameters.

    Args:
        kafka: initialized KafkaCommunicator object to preform actions with.
        demisto_params: The demisto parameters.


    Return 'ok' if everything went well, raise relevant exception otherwise
    """
    kafka.test_connection(log_stream=log_stream)
    if demisto_params.get("isFetch", False):
        check_params(
            kafka=kafka,
            topic=demisto_params.get("topic", None),
            partitions=handle_empty(argToList(demisto_params.get("partition", None)), None),
            offset=handle_empty(demisto_params.get("offset", "earliest"), "earliest"),
        )
    return "ok"


def print_topics(kafka: KafkaCommunicator, demisto_args: dict) -> Union[CommandResults, str]:
    """Print available topics in Broker

    Args:
        kafka: initialized KafkaCommunicator object to preform actions with.
        demisto_args: The demisto command arguments.

    Return CommandResults with the detailed topics, 'No topics found.' if no topics were found.
    """
    include_offsets = argToBoolean(demisto_args.get("include_offsets", "true"))
    kafka_topics = kafka.get_topics().values()
    if kafka_topics:
        topics = []
        for topic in kafka_topics:
            partitions = []
            for partition in topic.partitions.values():
                partition_output = {"ID": partition.id}
                if include_offsets:
                    try:
                        partition_output["EarliestOffset"], partition_output["OldestOffset"] = kafka.get_partition_offsets(
                            topic=topic.topic, partition=partition.id
                        )
                    except KafkaException as e:
                        # Sometimes listing topics can return uninitialized partitions.
                        # If that's the case, ignore them and continue.
                        if "Unknown partition" not in str(e):
                            raise e
                partitions.append(partition_output)

            topics.append({"Name": topic.topic, "Partitions": partitions})

        readable_output = tableToMarkdown("Kafka Topics", topics)

        return CommandResults(
            outputs_prefix="Kafka.Topic",
            outputs_key_field="Name",
            outputs=topics,
            readable_output=readable_output,
        )

    else:
        return "No topics found."


def produce_message(kafka: KafkaCommunicator, demisto_args: dict) -> None:
    """Producing message to kafka topic

    Args:
        kafka: initialized KafkaCommunicator object to preform actions with.
        demisto_args: The demisto command arguments.

    kafka.delivery_report is called when producing is done and it returns the relevant results.
    """
    topic = demisto_args.get("topic")
    value = demisto_args.get("value")
    value_schema_type = demisto_args.get("value_schema_type")
    value_schema_str = demisto_args.get("value_schema_str")
    value_schema_subject_name = demisto_args.get("value_schema_subject_name")
    partition_arg = demisto_args.get("partitioning_key")

    partition = -1
    if partition_arg is not None and str(partition_arg).isdigit():
        partition = int(partition_arg)

    try:
        kafka.produce(
            value=str(value),
            topic=str(topic),
            value_schema_type=value_schema_type,
            value_schema_str=value_schema_str,
            value_schema_subject_name=value_schema_subject_name,
            partition=partition,
        )
    except Exception as e:
        if "Topic authorization failed" in str(e):
            raise DemistoException(
                f"Error: {e!s}\n"
                "Check if you have permission to produce messages."
                "Your access might be restricted to consumer-only."
            )
        else:
            raise DemistoException(e)


def consume_message(kafka: KafkaCommunicator, demisto_args: dict) -> CommandResults | str:
    """Consume one message from topic

    Args:
        kafka: initialized KafkaCommunicator object to preform actions with.
        demisto_args: The demisto command arguments.

    Return CommandResults with the relevant message from kafka.
    """
    topic = str(demisto_args.get("topic"))
    partition = int(demisto_args.get("partition", -1))
    offset = demisto_args.get("offset", "0")

    message = kafka.consume(
        float(demisto_args.get("poll_timeout") or kafka.POLL_TIMEOUT), topic=topic, partition=partition, offset=offset
    )
    if not message:
        return "No message was consumed."
    else:
        message_value = message.value()
        if "Group authorization failed" in message_value.decode("utf-8"):
            raise DemistoException(f"{message_value} Make sure you configured the right Consumer group ID.")
        readable_output = tableToMarkdown(
            f"Message consumed from topic {topic}", [{"Offset": message.offset(), "Message": message_value.decode("utf-8")}]
        )
        content = {"Name": topic, "Message": {"Value": message_value.decode("utf-8"), "Offset": message.offset()}}

        return CommandResults(
            outputs=content, readable_output=readable_output, outputs_key_field="Name", outputs_prefix="Kafka.Topic"
        )


def fetch_partitions(kafka: KafkaCommunicator, demisto_args: dict) -> CommandResults:
    """Get available partitions in a given topic

    Args:
        kafka: initialized KafkaCommunicator object to preform actions with.
        demisto_args: The demisto command arguments.

    Return CommandResults with the relevant partitions.
    """
    topic = demisto_args.get("topic")
    kafka_topics = kafka.get_topics()
    if topic in kafka_topics:
        kafka_topic = kafka_topics[topic]
        partition_objects = kafka_topic.partitions.values()
        partitions = [partition.id for partition in partition_objects]

        readable_output = tableToMarkdown(name=f"Available partitions for topic '{topic}'", t=partitions, headers="Partitions")
        return CommandResults(
            outputs_prefix="Kafka.Topic",
            outputs_key_field="Name",
            outputs={"Name": topic, "Partition": partitions},
            readable_output=readable_output,
        )
    else:
        raise DemistoException(f"Topic {topic} was not found in Kafka")


def handle_empty(value: Any, default_value: Any) -> Any:
    """If value is empty return default_value."""
    if not value:
        return default_value
    return value


def check_params(
    kafka: KafkaCommunicator,
    topic: str,
    partitions: Optional[list] = None,
    offset: Optional[str] = None,
    consumer: bool = False,
    check_offset: bool = True,
) -> bool:
    """Check that partitions exist in topic and that offset matches the available ones.

    Args:
        kafka: Initialized KafkaCommunicator object to preform actions with.
        topic: The topic
        partitions: list of partitions
        offset: The offset
        consumer: consumer: Whether to get topics using the consumer's configuration (True) or producer's (False).

    Return True if everything is valid, raise relevant exception otherwise.
    """
    checkable_offset = False
    numerical_offset = 0
    topics = kafka.get_topics(consumer=consumer)
    if topic not in topics:
        raise DemistoException(f"Did not find topic {topic} in kafka topics.")

    if offset and str(offset).lower() not in SUPPORTED_GENERAL_OFFSETS:
        if offset.isdigit():
            numerical_offset = int(offset)
            checkable_offset = True
        else:
            raise DemistoException(f"Offset {offset} is not in supported format.")

    if partitions:
        topic_metadata = topics[topic]
        available_partitions = topic_metadata.partitions.values()
        available_partitions_ids = [available_partition.id for available_partition in available_partitions]
        for partition in partitions:
            if int(partition) not in available_partitions_ids:
                raise DemistoException(
                    f"Partition {partition} is not assigned to kafka topic {topic} available {available_partitions_ids}."
                )
            if check_offset and checkable_offset:
                earliest_offset, oldest_offset = kafka.get_partition_offsets(topic=topic, partition=int(partition))
                if numerical_offset < int(earliest_offset) or numerical_offset >= int(oldest_offset):
                    raise DemistoException(
                        f"Error checking params: Offset {numerical_offset} for topic {topic} and "
                        f"partition {partition} is out of bounds [{earliest_offset}, "
                        f"{oldest_offset})"
                    )

    return True


def get_topic_partition_if_relevant(kafka: KafkaCommunicator, topic: str, partition: str, specific_offset: str | int) -> list:
    """Return the TopicPartition if topic, partition and specific_offset are valid otherwise return []

    Args:
        kafka: The KafkaCommunicator object
        topic: The topic of the TopicPartition to retrieve
        partition: The partition of the TopicPartition to retrieve
        specific_offset: The offset of the TopicPartition to retrieve

    """
    demisto.debug(f"Getting last offset for partition {partition}, specific offset is {specific_offset}\n")
    add_topic_partition = True
    if isinstance(specific_offset, int):
        specific_offset += 1
        earliest_offset, latest_offset = kafka.get_partition_offsets(topic=topic, partition=int(partition))

        if specific_offset >= latest_offset or specific_offset < earliest_offset:
            add_topic_partition = False
            demisto.debug(
                f"Skipping partition {partition}, due to specific offset mismatch: "
                f"{specific_offset} not in [{earliest_offset}, {latest_offset}) \n"
            )

    if add_topic_partition:
        return kafka.get_topic_partitions(topic=topic, partition=int(partition), offset=specific_offset, consumer=True)
    return []


def get_fetch_topic_partitions(
    kafka: KafkaCommunicator, topic: str, offset: str | int, last_fetched_offsets: dict
) -> List[TopicPartition]:
    """Get topic partitions for fetching incidents without a specified partitions.

    If fetched from those partitions before chose offset accordingly.

    Args:
        kafka: The KafkaCommunicator object
        topic: The topic of the TopicPartitions to retrieve
        offset: The general offset to start fetching from
        last_fetched_offsets: The dictionary with the last fetched offsets

    Return a list of topic partitions
    """
    demisto.debug(f"Getting all topic partitions for topic {topic} and offset {offset}")
    all_topic_partitions = kafka.get_topic_partitions(topic=topic, partition=-1, offset=offset, consumer=True)
    if not last_fetched_offsets:
        demisto.debug("Did not fetch from this topic previously, returning all available topic partitions")
        return all_topic_partitions

    topic_partitions_in_system = []

    demisto.debug("Going over last fetched offsets")
    for partition in last_fetched_offsets:
        specific_offset = last_fetched_offsets.get(partition, offset)
        topic_partitions_in_system += get_topic_partition_if_relevant(kafka, topic, partition, specific_offset)

        for topic_partition in all_topic_partitions:
            if topic_partition.partition == int(partition):
                demisto.debug(f"Updating topic {topic} and partition {partition} to fetch from previous offset {specific_offset}")
                all_topic_partitions.remove(topic_partition)

    return topic_partitions_in_system + all_topic_partitions


def fetch_incidents(kafka: KafkaCommunicator, demisto_params: dict) -> None:
    """Fetch incidents as kafka messages from a specific topic.

    Args:
        kafka: initialized KafkaCommunicator object to preform actions with.
        demisto_params: The demisto parameters.
    """
    topic = demisto_params.get("topic", "")
    partitions = handle_empty(argToList(demisto_params.get("partition", "")), [])
    offset = handle_empty(demisto_params.get("first_fetch", "earliest"), "earliest")
    message_max_bytes = int(handle_empty(demisto_params.get("max_bytes_per_message", 1048576), 1048576))
    max_messages = int(handle_empty(demisto_params.get("max_fetch", 50), 50))
    last_fetched_offsets = demisto.getLastRun().get("last_fetched_offsets", {})
    last_topic = demisto.getLastRun().get("last_topic", "")
    stop_consuming_upon_timeout = argToBoolean(demisto_params.get("stop_consuming_upon_timeout", False))
    poll_timeout = kafka.POLL_TIMEOUT_STOP_UPON_TIMEOUT if stop_consuming_upon_timeout else kafka.POLL_TIMEOUT
    demisto.debug(
        f"Starting fetch incidents with:\n last_topic: {last_topic}, "
        f"last_fetched_offsets: {last_fetched_offsets}, "
        f"topic: {topic}, partitions: {partitions}, offset: {offset}, "
        f"message_max_bytes: {message_max_bytes}, max_messages: {max_messages}\n"
    )
    incidents = []

    kafka.update_conf_for_fetch(message_max_bytes=message_max_bytes)

    kafka_consumer = kafka.get_kafka_consumer()
    demisto.debug("Checking params")
    check_params(kafka, topic, partitions, offset, True, False)

    if topic != last_topic:
        demisto.debug(
            f"Topic changed from {last_topic} to {topic}, resetting last fetched offsets from "
            f"{last_fetched_offsets} to empty dict."
        )
        last_fetched_offsets = {}

    if offset.isdigit():
        offset = int(offset)

    topic_partitions = []
    for partition in partitions:
        specific_offset = last_fetched_offsets.get(partition, offset) if partition in last_fetched_offsets else offset
        topic_partitions += get_topic_partition_if_relevant(kafka, topic, partition, specific_offset)

    if not partitions:
        if isinstance(offset, int):
            offset += 1
        demisto.debug(f"No partitions were set, getting all available partitions for topic {topic}")
        topic_partitions = get_fetch_topic_partitions(kafka, topic, offset, last_fetched_offsets)

    try:
        demisto.debug(f"The topic partitions assigned to the consumer are: {topic_partitions}")

        if topic_partitions:
            kafka_consumer.assign(topic_partitions)

            demisto.debug("Beginning to poll messages from kafka")
            num_polled_msg = 0
            for _ in range(max_messages):
                # Initial message consumption may take up to
                # `session.timeout.ms` for the consumer group to
                # rebalance and start consuming
                polled_msg = kafka_consumer.poll(poll_timeout)
                if polled_msg:
                    num_polled_msg += 1
                    demisto.debug(f"Received a message {num_polled_msg}# from Kafka.")
                    incidents.append(create_incident(message=polled_msg, topic=topic))
                    last_fetched_offsets[f"{polled_msg.partition()}"] = polled_msg.offset()
                elif stop_consuming_upon_timeout and (not polled_msg):
                    demisto.debug(
                        f"Didn't get a message after {poll_timeout} seconds"
                        f", stop_consuming_upon_timeout is true, break the loop. {num_polled_msg=}"
                    )
                    break

    finally:
        if kafka_consumer:
            kafka_consumer.close()

    last_run = {"last_fetched_offsets": last_fetched_offsets, "last_topic": topic}
    demisto.debug(f"Fetching finished, setting last run to {last_run}")
    demisto.setLastRun(last_run)

    demisto.incidents(incidents)


""" COMMANDS MANAGER / SWITCH PANEL """


@capture_logs
def commands_manager(
    kafka_kwargs: dict,
    demisto_params: dict,
    demisto_args: dict,  # pragma: no cover
    demisto_command: str,
    kafka_logger: Optional[logging.Logger] = None,
    log_stream: Optional[StringIO] = None,
) -> None:
    """Start command function according to demisto command."""

    kafka_kwargs["kafka_logger"] = kafka_logger
    kafka = KafkaCommunicator(**kafka_kwargs)

    try:
        if demisto_command == "test-module":
            return_results(command_test_module(kafka, demisto_params, log_stream))
        elif demisto_command == "kafka-print-topics":
            return_results(print_topics(kafka, demisto_args))
        elif demisto_command == "kafka-publish-msg":
            produce_message(kafka, demisto_args)
        elif demisto_command == "kafka-consume-msg":
            return_results(consume_message(kafka, demisto_args))
        elif demisto_command == "kafka-fetch-partitions":
            return_results(fetch_partitions(kafka, demisto_args))
        elif demisto_command == "fetch-incidents":
            fetch_incidents(kafka, demisto_params)
        else:
            raise NotImplementedError(f"Command {demisto_command} not found in command list")
    finally:
        if kafka.ca_path and os.path.isfile(kafka.ca_path):
            os.remove(os.path.abspath(kafka.ca_path))
        if kafka.client_cert_path and os.path.isfile(kafka.client_cert_path):
            os.remove(os.path.abspath(kafka.client_cert_path))
        if kafka.client_key_path and os.path.isfile(kafka.client_key_path):
            os.remove(os.path.abspath(kafka.client_key_path))


def main():  # pragma: no cover
    demisto_command = demisto.command()
    demisto_params = demisto.params()
    demisto_args = demisto.args()
    demisto.debug(f"Command being called is {demisto_command}")
    brokers = demisto_params.get("brokers")
    group_id = demisto_params.get("group_id", "xsoar_group")
    offset = handle_empty(demisto_params.get("offset", "earliest"), "earliest")
    trust_any_cert = demisto_params.get("insecure", False)

    # Should we use SSL
    use_ssl = demisto_params.get("use_ssl", False)
    # Should we use SASL (with SSL and PLAIN)
    use_sasl = demisto_params.get("use_sasl", False)

    ca_cert = demisto_params.get("ca_cert", None)
    client_cert = demisto_params.get("client_cert", None)
    client_cert_key = demisto_params.get("client_cert_key", None)
    ssl_password = demisto_params.get("additional_password", None)
    plain_username = demisto_params.get("credentials", {}).get("identifier")
    plain_password = demisto_params.get("credentials", {}).get("password")
    schema_registry_url = demisto_params.get("schema_registry_url", None)
    schema_registry_username = demisto_params.get("schema_registry_credentials", {}).get("identifier", None)
    schema_registry_password = demisto_params.get("schema_registry_credentials", {}).get("password", None)
    validate_params(
        use_ssl=use_ssl,
        use_sasl=use_sasl,
        trust_any_cert=trust_any_cert,
        plain_username=plain_username,
        plain_password=plain_password,
        brokers=brokers,
        ca_cert=ca_cert,
        client_cert=client_cert,
        client_cert_key=client_cert_key,
    )

    kafka_kwargs = {
        "use_ssl": use_ssl,
        "brokers": brokers,
        "ca_cert": ca_cert,
        "offset": offset,
        "use_sasl": use_sasl,
        "group_id": group_id,
        "trust_any_cert": trust_any_cert,
        "client_cert": client_cert,
        "client_cert_key": client_cert_key,
        "plain_username": plain_username,
        "plain_password": plain_password,
        "schema_registry_url": schema_registry_url,
        "schema_registry_username": schema_registry_username,
        "schema_registry_password": schema_registry_password,
    }
    if ssl_password:
        kafka_kwargs["ssl_password"] = ssl_password

    try:
        commands_manager(kafka_kwargs, demisto_params, demisto_args, demisto_command)

    except Exception as e:
        debug_log = ""
        stacktrace = traceback.format_exc()
        if stacktrace:
            debug_log += f"Debug logs:\nFull stacktrace:\n\n{stacktrace}"
        return_error(f"{e!s}\n\n{debug_log}")


if __name__ in ("__main__", "__builtin__", "builtins"):  # pragma: no cover
    main()