/* -*- Mode: C; tab-width: 4; indent-tabs-mode: nil; c-basic-offset: 4 -*- */

/*  Fluent Bit

 *  ==========

 *  Copyright (C) 2019-2021 The Fluent Bit Authors

 *  Copyright (C) 2015-2018 Treasure Data Inc.

 *

 *  Licensed under the Apache License, Version 2.0 (the "License");

 *  you may not use this file except in compliance with the License.

 *  You may obtain a copy of the License at

 *

 *      http://www.apache.org/licenses/LICENSE-2.0

 *

 *  Unless required by applicable law or agreed to in writing, software

 *  distributed under the License is distributed on an "AS IS" BASIS,

 *  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

 *  See the License for the specific language governing permissions and

 *  limitations under the License.

 */

#include <fluent-bit/flb_input_plugin.h>

#include <fluent-bit/flb_pack.h>

#include <fluent-bit/flb_engine.h>

#include <fluent-bit/flb_time.h>

#include <fluent-bit/flb_parser.h>

#include <fluent-bit/flb_error.h>

#include <fluent-bit/flb_utils.h>

#include <fluent-bit/aws/flb_aws_msk_iam.h>

#include <mpack/mpack.h>

#include <stddef.h>

#include <stdio.h>

#include "fluent-bit/flb_input.h"

#include "fluent-bit/flb_kafka.h"

#include "fluent-bit/flb_mem.h"

#include "in_kafka.h"

#include "rdkafka.h"

static int try_json(struct flb_log_event_encoder *log_encoder,

                    rd_kafka_message_t *rkm)

{

    int root_type;

    char *buf = NULL;

    size_t bufsize;

    int ret;

    ret = flb_pack_json(rkm->payload, rkm->len, &buf, &bufsize, &root_type, NULL);

    if (ret) {

        if (buf) {

            flb_free(buf);

        }

        return ret;

    }

    flb_log_event_encoder_append_body_raw_msgpack(log_encoder, buf, bufsize);

    flb_free(buf);

    return 0;

}

static int process_message(struct flb_in_kafka_config *ctx,

                           rd_kafka_message_t *rkm)

{

    struct flb_log_event_encoder *log_encoder = ctx->log_encoder;

    int ret;

    ret = flb_log_event_encoder_begin_record(log_encoder);

    if (ret == FLB_EVENT_ENCODER_SUCCESS) {

        ret = flb_log_event_encoder_set_current_timestamp(log_encoder);

    }

    if (ret == FLB_EVENT_ENCODER_SUCCESS) {

        ret = flb_log_event_encoder_append_body_cstring(log_encoder, "topic");

    }

    if (ret == FLB_EVENT_ENCODER_SUCCESS) {

        if (rkm->rkt) {

            ret = flb_log_event_encoder_append_body_cstring(log_encoder,

                                                            (char *) rd_kafka_topic_name(rkm->rkt));

        }

        else {

            ret = flb_log_event_encoder_append_body_null(log_encoder);

        }

    }

    if (ret == FLB_EVENT_ENCODER_SUCCESS) {

        ret = flb_log_event_encoder_append_body_values(log_encoder,

                                                       FLB_LOG_EVENT_CSTRING_VALUE("partition"),

                                                       FLB_LOG_EVENT_INT32_VALUE(rkm->partition));

    }

    if (ret == FLB_EVENT_ENCODER_SUCCESS) {

        ret = flb_log_event_encoder_append_body_values(log_encoder,

                                                       FLB_LOG_EVENT_CSTRING_VALUE("offset"),

                                                       FLB_LOG_EVENT_INT64_VALUE(rkm->offset));

    }

    if (ret == FLB_EVENT_ENCODER_SUCCESS) {

        ret = flb_log_event_encoder_append_body_cstring(log_encoder, "error");

    }

    if (ret == FLB_EVENT_ENCODER_SUCCESS) {

        if (rkm->err) {

            ret = flb_log_event_encoder_append_body_cstring(log_encoder,

                                                            (char *) rd_kafka_message_errstr(rkm));

        }

        else {

            ret = flb_log_event_encoder_append_body_null(log_encoder);

        }

    }

    if (ret == FLB_EVENT_ENCODER_SUCCESS) {

        ret = flb_log_event_encoder_append_body_cstring(log_encoder, "key");

    }

    if (ret == FLB_EVENT_ENCODER_SUCCESS) {

        if (rkm->key) {

            ret = flb_log_event_encoder_append_body_string(log_encoder,

                                                           rkm->key,

                                                           rkm->key_len);

        }

        else {

            ret = flb_log_event_encoder_append_body_null(log_encoder);

        }

    }

    if (ret == FLB_EVENT_ENCODER_SUCCESS) {

        ret = flb_log_event_encoder_append_body_cstring(log_encoder, "payload");

    }

    if (ret == FLB_EVENT_ENCODER_SUCCESS) {

        if (rkm->payload) {

            if (ctx->format != FLB_IN_KAFKA_FORMAT_JSON ||

                    try_json(log_encoder, rkm)) {

                ret = flb_log_event_encoder_append_body_string(log_encoder,

                                                               rkm->payload,

                                                               rkm->len);

            }

        }

        else {

            ret = flb_log_event_encoder_append_body_null(log_encoder);

        }

    }

    if (ret == FLB_EVENT_ENCODER_SUCCESS) {

        ret = flb_log_event_encoder_commit_record(log_encoder);

    }

    if (ret != FLB_EVENT_ENCODER_SUCCESS) {

        flb_log_event_encoder_rollback_record(log_encoder);

    }

    return ret;

}

static int in_kafka_collect(struct flb_input_instance *ins,

                            struct flb_config *config, void *in_context)

{

    int ret;

    int append_ret;

    struct flb_in_kafka_config *ctx = in_context;

    rd_kafka_message_t *rkm;

    ret = FLB_EVENT_ENCODER_SUCCESS;

    while (ret == FLB_EVENT_ENCODER_SUCCESS) {

        /* Set the Kafka poll timeout based on execution mode:

         *

         * a) Running in the main event loop (non-threaded):

         *    - Use a minimal timeout to avoid blocking other inputs.

         *

         * b) Running in a dedicated thread:

         *    - Optimize for throughput by allowing Kafka's internal batching.

         *    - Align with 'fetch.wait.max.ms' (default: 500ms) to maximize batch efficiency.

         *    - Set timeout slightly higher than 'fetch.wait.max.ms' (e.g., 1.5x - 2x) to

         *      ensure it does not interfere with Kafka's fetch behavior, while still

         *      keeping the consumer responsive.

         */

        if (ctx->ins->flags & FLB_INPUT_THREADED) {

            /* Threaded mode: Optimize for batch processing and efficiency */

            rkm = rd_kafka_consumer_poll(ctx->kafka.rk, ctx->poll_timeout_ms);

        } else {

            /* Main event loop: Minimize delay for non-blocking execution */

            rkm = rd_kafka_consumer_poll(ctx->kafka.rk, 1);

        }

        if (!rkm) {

            break;

        }

        if (rkm->err) {

            flb_plg_warn(ins, "consumer error: %s\n",

                         rd_kafka_message_errstr(rkm));

            rd_kafka_message_destroy(rkm);

            continue;

        }

        flb_plg_debug(ins, "kafka message received");

        ret = process_message(ctx, rkm);

        rd_kafka_message_destroy(rkm);

        /* Break from the loop when reaching the limit of polling if available */

        if (ctx->polling_threshold != FLB_IN_KAFKA_UNLIMITED &&

            ctx->log_encoder->output_length > ctx->polling_threshold + 512) {

            break;

        }

    }

    if (ret == FLB_EVENT_ENCODER_SUCCESS) {

        if (ctx->log_encoder->output_length > 0) {

            append_ret = flb_input_log_append(ins, NULL, 0,

                                              ctx->log_encoder->output_buffer,

                                              ctx->log_encoder->output_length);

            if (append_ret == 0) {

                if (!ctx->enable_auto_commit) {

                    rd_kafka_commit(ctx->kafka.rk, NULL, 0);

                }

                ret = 0;

            }

            else {

                flb_plg_error(ins, "failed to append records");

                ret = -1;

            }

        }

        else {

            ret = 0;

        }

    }

    else {

        flb_plg_error(ins, "Error encoding record : %d", ret);

        ret = -1;

    }

    flb_log_event_encoder_reset(ctx->log_encoder);

    return ret;

}

/* Initialize plugin */

static int in_kafka_init(struct flb_input_instance *ins,

                         struct flb_config *config, void *data)

{

    int ret;

    const char *conf;

    struct flb_in_kafka_config *ctx;

    rd_kafka_conf_t *kafka_conf = NULL;

    rd_kafka_topic_partition_list_t *kafka_topics = NULL;

    rd_kafka_resp_err_t err = RD_KAFKA_RESP_ERR_NO_ERROR;

    rd_kafka_conf_res_t res;

    char errstr[512];

    (void) data;

    char conf_val[16];

    /* Allocate space for the configuration context */

    ctx = flb_calloc(1, sizeof(struct flb_in_kafka_config));

    if (!ctx) {

        return -1;

    }

    ctx->ins = ins;

    ret = flb_input_config_map_set(ins, (void*) ctx);

    if (ret == -1) {

        flb_plg_error(ins, "unable to load configuration.");

        flb_free(ctx);

        return -1;

    }

#ifdef FLB_HAVE_AWS_MSK_IAM

    /*

     * When MSK IAM auth is enabled, default the required

     * security settings so users don't need to specify them.

     */

    if (ctx->aws_msk_iam && ctx->aws_msk_iam_cluster_arn) {

        conf = flb_input_get_property("rdkafka.security.protocol", ins);

        if (!conf) {

            flb_input_set_property(ins, "rdkafka.security.protocol", "SASL_SSL");

        }

        conf = flb_input_get_property("rdkafka.sasl.mechanism", ins);

        if (!conf) {

            flb_input_set_property(ins, "rdkafka.sasl.mechanism", "OAUTHBEARER");

            ctx->sasl_mechanism = flb_sds_create("OAUTHBEARER");

        }

        else {

            ctx->sasl_mechanism = flb_sds_create(conf);

            flb_plg_info(ins, "SASL mechanism configured: %s", ctx->sasl_mechanism);

        }

    }

    else {

#endif

        /* Retrieve SASL mechanism if configured */

        conf = flb_input_get_property("rdkafka.sasl.mechanism", ins);

        if (conf) {

            ctx->sasl_mechanism = flb_sds_create(conf);

            flb_plg_info(ins, "SASL mechanism configured: %s", ctx->sasl_mechanism);

        }

#ifdef FLB_HAVE_AWS_MSK_IAM

    }

#endif

    kafka_conf = flb_kafka_conf_create(&ctx->kafka, &ins->properties, 1);

    if (!kafka_conf) {

        flb_plg_error(ins, "Could not initialize kafka config object");

        goto init_error;

    }

    /* Set enable.auto.commit based on plugin's enable_auto_commit setting */

    res = rd_kafka_conf_set(kafka_conf, "enable.auto.commit",

                           ctx->enable_auto_commit ? "true" : "false",

                           errstr, sizeof(errstr));

    if (res != RD_KAFKA_CONF_OK) {

        flb_plg_error(ins, "Failed to set enable.auto.commit: %s", errstr);

        goto init_error;

    }

    if (ctx->buffer_max_size > 0) {

        ctx->polling_threshold = ctx->buffer_max_size;

        snprintf(conf_val, sizeof(conf_val), "%zu", ctx->polling_threshold - 512);

        res = rd_kafka_conf_set(kafka_conf, "fetch.max.bytes", conf_val,

                                errstr, sizeof(errstr));

        if (res != RD_KAFKA_CONF_OK) {

            flb_plg_error(ins, "Failed to set up fetch.max.bytes: %s, val = %s",

                          rd_kafka_err2str(err), conf_val);

            goto init_error;

        }

        snprintf(conf_val, sizeof(conf_val), "%zu", ctx->polling_threshold);

        res = rd_kafka_conf_set(kafka_conf, "receive.message.max.bytes", conf_val,

                                errstr, sizeof(errstr));

        if (res != RD_KAFKA_CONF_OK) {

            flb_plg_error(ins, "Failed to set up receive.message.max.bytes: %s, val = %s",

                          rd_kafka_err2str(err), conf_val);

            goto init_error;

        }

    }

    else {

        ctx->polling_threshold = FLB_IN_KAFKA_UNLIMITED;

    }

    ctx->opaque = flb_kafka_opaque_create();

    if (!ctx->opaque) {

        flb_plg_error(ins, "failed to create kafka opaque context");

        goto init_error;

    }

    flb_kafka_opaque_set(ctx->opaque, ctx, NULL);

    rd_kafka_conf_set_opaque(kafka_conf, ctx->opaque);

#ifdef FLB_HAVE_AWS_MSK_IAM

    if (ctx->aws_msk_iam && ctx->aws_msk_iam_cluster_arn && ctx->sasl_mechanism &&

        strcasecmp(ctx->sasl_mechanism, "OAUTHBEARER") == 0) {

        flb_plg_info(ins, "registering MSK IAM authentication with cluster ARN: %s",

                     ctx->aws_msk_iam_cluster_arn);

        ctx->msk_iam = flb_aws_msk_iam_register_oauth_cb(config,

                                                         kafka_conf,

                                                         ctx->aws_msk_iam_cluster_arn,

                                                         ctx->opaque);

        if (!ctx->msk_iam) {

            flb_plg_error(ins, "failed to setup MSK IAM authentication");

        }

        else {

            res = rd_kafka_conf_set(kafka_conf, "sasl.oauthbearer.config",

                                    "principal=admin", errstr, sizeof(errstr));

            if (res != RD_KAFKA_CONF_OK) {

                flb_plg_error(ins,

                             "failed to set sasl.oauthbearer.config: %s",

                             errstr);

            }

        }

    }

#endif

    ctx->kafka.rk = rd_kafka_new(RD_KAFKA_CONSUMER, kafka_conf, errstr, sizeof(errstr));

    /* Create Kafka consumer handle */

    if (!ctx->kafka.rk) {

        flb_plg_error(ins, "Failed to create new consumer: %s", errstr);

        goto init_error;

    }

    /* Trigger initial token refresh for OAUTHBEARER */

    rd_kafka_poll(ctx->kafka.rk, 0);

    conf = flb_input_get_property("topics", ins);

    if (!conf) {

        flb_plg_error(ins, "config: no topics specified");

        goto init_error;

    }

    kafka_topics = flb_kafka_parse_topics(conf);

    if (!kafka_topics) {

        flb_plg_error(ins, "Failed to parse topic list");

        goto init_error;

    }

    if (strcasecmp(ctx->format_str, "none") == 0) {

        ctx->format = FLB_IN_KAFKA_FORMAT_NONE;

    }

    else if (strcasecmp(ctx->format_str, "json") == 0) {

        ctx->format = FLB_IN_KAFKA_FORMAT_JSON;

    }

    else {

        flb_plg_error(ins, "config: invalid format \"%s\"", ctx->format_str);

        goto init_error;

    }

    if ((err = rd_kafka_subscribe(ctx->kafka.rk, kafka_topics))) {

        flb_plg_error(ins, "Failed to start consuming topics: %s", rd_kafka_err2str(err));

        goto init_error;

    }

    rd_kafka_topic_partition_list_destroy(kafka_topics);

    kafka_topics = NULL;

    /* Set the context */

    flb_input_set_context(ins, ctx);

    /* Collect upon data available on the pipe read fd */

    int poll_seconds = ctx->poll_ms / 1000;

    int poll_milliseconds = ctx->poll_ms % 1000;

    ret = flb_input_set_collector_time(ins,

                                       in_kafka_collect,

                                       poll_seconds, poll_milliseconds * 1e6,

                                       config);

    if (ret) {

        flb_plg_error(ctx->ins, "could not set collector for kafka input plugin");

        goto init_error;

    }

    ctx->coll_fd = ret;

    ctx->log_encoder = flb_log_event_encoder_create(FLB_LOG_EVENT_FORMAT_DEFAULT);

    if (ctx->log_encoder == NULL) {

        flb_plg_error(ins, "could not initialize log encoder");

        goto init_error;

    }

    return 0;

init_error:

    if (kafka_topics) {

        rd_kafka_topic_partition_list_destroy(kafka_topics);

    }

    if (ctx->kafka.rk) {

        rd_kafka_consumer_close(ctx->kafka.rk);

        rd_kafka_destroy(ctx->kafka.rk);

    }

    if (ctx->opaque) {

        flb_kafka_opaque_destroy(ctx->opaque);

    }

    else if (kafka_conf) {

        /* conf is already destroyed when rd_kafka is initialized */

        rd_kafka_conf_destroy(kafka_conf);

    }

    flb_sds_destroy(ctx->sasl_mechanism);

    flb_free(ctx);

    return -1;

}

static void in_kafka_pause(void *data, struct flb_config *config)

{

    struct flb_in_kafka_config *ctx = data;

    flb_input_collector_pause(ctx->coll_fd, ctx->ins);

}

static void in_kafka_resume(void *data, struct flb_config *config)

{

    struct flb_in_kafka_config *ctx = data;

    flb_input_collector_resume(ctx->coll_fd, ctx->ins);

}

/* Cleanup serial input */

static int in_kafka_exit(void *in_context, struct flb_config *config)

{

    struct flb_in_kafka_config *ctx;

    if (!in_context) {

        return 0;

    }

    ctx = in_context;

    if (ctx->kafka.rk) {

        rd_kafka_consumer_close(ctx->kafka.rk);

        rd_kafka_destroy(ctx->kafka.rk);

    }

    flb_free(ctx->kafka.brokers);

    if (ctx->log_encoder){

        flb_log_event_encoder_destroy(ctx->log_encoder);

    }

#ifdef FLB_HAVE_AWS_MSK_IAM

    if (ctx->msk_iam) {

        flb_aws_msk_iam_destroy(ctx->msk_iam);

    }

#endif

    if (ctx->opaque) {

        flb_kafka_opaque_destroy(ctx->opaque);

    }

    flb_sds_destroy(ctx->sasl_mechanism);

    flb_free(ctx);

    return 0;

}

static struct flb_config_map config_map[] = {

   {

    FLB_CONFIG_MAP_INT, "poll_ms", FLB_IN_KAFKA_DEFAULT_POLL_MS,

    0, FLB_TRUE, offsetof(struct flb_in_kafka_config, poll_ms),

    "Interval in milliseconds to check for new messages."

   },

   {

    FLB_CONFIG_MAP_STR, "topics", (char *)NULL,

    0, FLB_FALSE, 0,

    "Set the kafka topics, delimited by commas."

   },

   {

    FLB_CONFIG_MAP_STR, "format", FLB_IN_KAFKA_DEFAULT_FORMAT,

    0, FLB_TRUE, offsetof(struct flb_in_kafka_config, format_str),

    "Set the data format which will be used for parsing records."

   },

   {

    FLB_CONFIG_MAP_STR, "brokers", (char *)NULL,

    0, FLB_FALSE, 0,

    "Set the kafka brokers, delimited by commas."

   },

   {

    FLB_CONFIG_MAP_STR, "client_id", (char *)NULL,

    0, FLB_FALSE, 0,

    "Set the kafka client_id."

   },

   {

    FLB_CONFIG_MAP_STR, "group_id", (char *)NULL,

    0, FLB_FALSE, 0,

    "Set the kafka group_id."

   },

   {

    FLB_CONFIG_MAP_STR_PREFIX, "rdkafka.", NULL,

    /* FLB_CONFIG_MAP_MULT, FLB_TRUE, offsetof(struct flb_in_kafka_config, rdkafka_opts), */

    0,  FLB_FALSE, 0,

    "Set the librdkafka options"

   },

   {

    FLB_CONFIG_MAP_SIZE, "buffer_max_size", FLB_IN_KAFKA_BUFFER_MAX_SIZE,

    0, FLB_TRUE, offsetof(struct flb_in_kafka_config, buffer_max_size),

    "Set the maximum size of chunk"

   },

   {

    FLB_CONFIG_MAP_INT, "poll_timeout_ms", "1",

    0, FLB_TRUE, offsetof(struct flb_in_kafka_config, poll_timeout_ms),

    "Set the timeout in milliseconds for Kafka consumer poll operations. "

    "This option only takes effect when running in a dedicated thread (i.e., when 'threaded' is enabled). "

    "Using a higher timeout (e.g., 1.5x - 2x 'rdkafka.fetch.wait.max.ms') "

    "can improve efficiency by leveraging Kafka's batching mechanism."

  },

  {

    FLB_CONFIG_MAP_BOOL, "enable_auto_commit", FLB_IN_KAFKA_ENABLE_AUTO_COMMIT,

    0, FLB_TRUE, offsetof(struct flb_in_kafka_config, enable_auto_commit),

    "Rely on kafka auto-commit and commit messages in batches"

  },

#ifdef FLB_HAVE_AWS_MSK_IAM

  {

   FLB_CONFIG_MAP_STR, "aws_msk_iam_cluster_arn", (char *)NULL,

   0, FLB_TRUE, offsetof(struct flb_in_kafka_config, aws_msk_iam_cluster_arn),

   "ARN of the MSK cluster when using AWS IAM authentication"

  },

  {

    FLB_CONFIG_MAP_BOOL, "aws_msk_iam", "false",

    0, FLB_TRUE, offsetof(struct flb_in_kafka_config, aws_msk_iam),

    "Enable AWS MSK IAM authentication"

  },

#endif

  /* EOF */

  {0}

};

/* Plugin reference */

struct flb_input_plugin in_kafka_plugin = {

    .name         = "kafka",

    .description  = "Kafka consumer input plugin",

    .cb_init      = in_kafka_init,

    .cb_pre_run   = NULL,

    .cb_collect   = in_kafka_collect,

    .cb_flush_buf = NULL,

    .cb_pause     = in_kafka_pause,

    .cb_resume    = in_kafka_resume,

    .cb_exit      = in_kafka_exit,

    .config_map   = config_map

};