/*
 * Creator: Naman Dixit
 * Notice: © Copyright 2020 Naman Dixit
 */

#define logMessage(s, ...) say(s "\n", ##__VA_ARGS__)
#define logError(s, ...)   err(s "\n", ##__VA_ARGS__)

#include "nlib/nlib.h"

#undef internal_malloc
#undef internal_free

#include <sys/socket.h>
#include <sys/epoll.h>
#include <netdb.h>
#include <unistd.h>
#include <fcntl.h>
#include <sys/time.h>
#include <arpa/inet.h>
#include <assert.h>
#include <signal.h>
#include <librdkafka/rdkafka.h>

#include "kafka.h"
#include "time.c"
#include "conf.c"
#include "command.c"
#include "grunt_manager.c"

# if defined(COMPILER_CLANG)
#  pragma clang diagnostic push
#   pragma clang diagnostic ignored "-Wpadded"
#   pragma clang diagnostic ignored "-Wfloat-equal"
# endif
#include "cJSON/cJSON.h"
#include "cJSON/cJSON.c"
# if defined(COMPILER_CLANG)
#  pragma clang diagnostic pop
# endif

# if defined(COMPILER_CLANG)
#  pragma clang diagnostic push
#   pragma clang diagnostic ignored "-Wreserved-id-macro"
#   pragma clang diagnostic ignored "-Wcast-qual"
# endif
#include "inih/ini.h"
#include "inih/ini.c"
# if defined(COMPILER_CLANG)
#  pragma clang diagnostic pop
# endif

global_variable volatile sig_atomic_t global_keep_running = 1;

internal_function
void signalHandlerSIGINT (int _)
{
    (void)_;
    global_keep_running = 0;
}

Sint main (Sint argc, Char *argv[])
{
    unused_variable(argc);
    unused_variable(argv);

    signal(SIGINT, signalHandlerSIGINT);

    { // Init cJSON
        cJSON_Hooks hook = {.malloc_fn = dlmalloc,
                            .free_fn = dlfree};
        cJSON_InitHooks(&hook);
    }

    Configuration conf = {0};
    { // Default config values
        conf.kafka_address = "10.129.6.5:9092";

        conf.message_read_gap = 10;
        conf.grunt_time_to_die = 10000;
        conf.grunt_response_wait_time = 100;
    }

    if (ini_parse("config.ini", confCallback, &conf) < 0) {
        printf("Can't load 'config.ini'\n");
        return -1;
    }

    Command *commands = NULL;
    Grunt_Tracker gt = {0};

    Hash_Table grunt_survey_map = htCreate(0);

    Kafka kafka = {0};

    kafkaCreateWriter(&kafka, conf.kafka_address);

#define CREATE_TOPIC(s)                                 \
    do {                                                \
        if (kafkaCreateTopic(&kafka, s, 1, 1) == -1) {  \
            rd_kafka_destroy(kafka.writer);             \
            return -1;                                  \
        }                                               \
    } while (0)

    CREATE_TOPIC("REQUEST_DM_2_RM"); //
    CREATE_TOPIC("RESPONSE_RM_2_DM");
    CREATE_TOPIC("REQUEST_RM_2_RD");
    CREATE_TOPIC("RESPONSE_RD_2_RM"); //
    CREATE_TOPIC("JOIN_RD_2_RM"); //
    CREATE_TOPIC("HEARTBEAT_RD_2_RM"); //
    CREATE_TOPIC("JOIN_ACK_RM_2_RD"); //
    CREATE_TOPIC("REJOIN_RM_2_RD"); //
    CREATE_TOPIC("LOG_COMMON"); //
#undef CREATE_TOPIC

    kafkaCreateReader(&kafka, conf.kafka_address);

    rd_kafka_topic_partition_list_t *kafka_reader_topics = rd_kafka_topic_partition_list_new(1);

    rd_kafka_topic_t *topic_req_dm2a = kafkaSubscribe(&kafka, kafka_reader_topics,
                                                      "REQUEST_DM_2_RM");
    rd_kafka_topic_t *topic_join_g2a = kafkaSubscribe(&kafka, kafka_reader_topics,
                                                      "JOIN_RD_2_RM");
    rd_kafka_topic_t *topic_res_g2a = kafkaSubscribe(&kafka, kafka_reader_topics,
                                                     "RESPONSE_RD_2_RM");
    rd_kafka_topic_t *topic_beat_g2a = kafkaSubscribe(&kafka, kafka_reader_topics,
                                                      "HEARTBEAT_RD_2_RM");
    rd_kafka_topic_t *topic_log = kafkaSubscribe(&kafka, kafka_reader_topics,
                                                 "LOG_COMMON");

    rd_kafka_resp_err_t kafka_reader_topics_err = rd_kafka_subscribe(kafka.reader,
                                                                     kafka_reader_topics);

    logMessage("Subscription finished\n");
    fflush(stdout);

    rd_kafka_topic_partition_list_destroy(kafka_reader_topics);

    if (kafka_reader_topics_err) {
        logError("Subscribe failed: %s\n",
                 rd_kafka_err2str(kafka_reader_topics_err));
        rd_kafka_destroy(kafka.reader);
        return -1;
    }

    Sint time_of_launch = (Sint)time(0);

    U64 time_passed_last = timeMilli();
    while (global_keep_running) {
        // NOTE(naman): Get the fd's that are ready
        rd_kafka_message_t *kafka_message_read = rd_kafka_consumer_poll(kafka.reader,
                                                                        conf.message_read_gap);

        U64 time_passed_now = timeMilli();
        U64 time_passed = time_passed_now - time_passed_last;
        time_passed_last = time_passed_now;

        for (Size j = 0; j < gt.map.slot_count; j++) {
            if (gt.map.keys[j] != 0) {
                gt.grunts[gt.map.values[j]].time_to_die -= time_passed;
            }
        }

        if (kafka_message_read != NULL) {
            if (kafka_message_read->err) {
                /* Consumer error: typically just informational. */
                logError("Consumer error: %s\n",
                         rd_kafka_message_errstr(kafka_message_read));
            } else {
                /* Proper message */
                /* fprintf(stderr, */
                /*         "Received message on %s [%d] " */
                /*         "at offset %"PRId64": \n%s\n", */
                /*         rd_kafka_topic_name(kafka_message_read->rkt), */
                /*         (int)kafka_message_read->partition, kafka_message_read->offset, */
                /*         cJSON_Print(cJSON_Parse((char *)kafka_message_read->payload))); */

                const char *json_error = NULL;
                cJSON *root = cJSON_ParseWithOpts(kafka_message_read->payload, &json_error, true);

                if ((cJSON_GetObjectItem(root, "timestamp") == NULL) ||
                    (cJSON_GetObjectItem(root, "timestamp")->valueint) < time_of_launch) {
                    logMessage("Ignoring : %s\n", kafka_message_read->payload);
                    cJSON_Delete(root);
                    rd_kafka_message_destroy(kafka_message_read);
                    continue;
                }

                if (kafka_message_read->rkt == topic_req_dm2a) {
                    Command c = {.kind = Command_RESPONSE_ARBITER_2_DM};
                    c.resource_id = strdup(cJSON_GetObjectItem(root, "resource_id")->valuestring);

                    // TODO(naman): Add any new resource fields here
                    Sint memory = cJSON_GetObjectItem(root, "memory")->valueint;
                    logMessage("Request DM2RM:\tid: %s = ([memory] = %d)",
                               c.resource_id, memory);

                    Char **grunt_ids = NULL;

                    for (Size j = 0; j < gt.map.slot_count; j++) {
                        if (gt.map.values[j] != 0) {
                            Grunt g = gt.grunts[gt.map.values[j]];
                            if (g.memory >= memory) {
                                c.kind = Command_RESPONSE_ARBITER_2_DM;
                                sbufAdd(grunt_ids, strdup(g.id));
                            }
                        }
                    }

                    if (c.kind == Command_REQUEST_ARBITER_2_GRUNT) {
                        c.req_a2g.memory = memory;
                        Grunt_Survey *gs = calloc(1, sizeof(*gs));
                        htInsert(&grunt_survey_map, hashString(c.resource_id), (Uptr)gs);

                        gs->milli_last = timeMilli();
                        gs->resource_id = c.resource_id;
                    } else if (c.kind == Command_RESPONSE_ARBITER_2_DM) {
                        c.res_a2d.grunt_ids = grunt_ids;
                        sbufAdd(commands, c);
                    }
                } else if (kafka_message_read->rkt == topic_join_g2a) {
                    Char *id = strdup(cJSON_GetObjectItem(root, "node_id")->valuestring);
                    Grunt grunt = {.id = strdup(id), .time_to_die = conf.grunt_time_to_die};

                    logMessage("Join G2A:\tid: %s", id);

                    if (htLookup(&gt.map, hashString(id)) == 0) {
                        gruntTrackBegin(&gt, grunt);
                    }

                    Command c = {.kind = Command_JOIN_ACK_ARBITER_2_GRUNT,
                                 .join_ack_a2g.grunt_id = strdup(id)};
                    sbufAdd(commands, c);
                } else if (kafka_message_read->rkt == topic_beat_g2a) {
                    Char *id = strdup(cJSON_GetObjectItem(root, "node_id")->valuestring);

                    U64 index = htLookup(&gt.map, hashString(id));
                    if (index != 0) { // Prevent any left over message
                        // TODO(naman): Add any new resource fields here
                        gt.grunts[index].time_to_die = conf.grunt_time_to_die;

                        gt.grunts[index].memory = cJSON_GetObjectItem(root, "memory")->valueint;
                        logMessage("Beat G2A:\tid: %s (Memory %d)", id, gt.grunts[index].memory);
                    } else {
                        if ((gt.grunts != NULL) && (gt.grunts[index].rejoin_asked != true)) {
                            gt.grunts[index].rejoin_asked = true;
                        }
                        Command c = {.kind = Command_REJOIN_ARBITER_2_GRUNT,
                                     .rejoin_a2g.grunt_id = id};
                        sbufAdd(commands, c);
                        logMessage("Beat G2A:\tid: %s (UNNOWN)", c.rejoin_a2g.grunt_id);
                    }
                } else if (kafka_message_read->rkt == topic_res_g2a) {
                    Char *node_id = cJSON_GetObjectItem(root, "node_id")->valuestring;
                    Char *resource_id = cJSON_GetObjectItem(root, "resource_id")->valuestring;
                    B32 success = (B32)cJSON_IsTrue(cJSON_GetObjectItem(root, "success"));

                    logMessage("Response G2A:\tid: %s (%s) = %s",
                               resource_id, node_id, success ? "succeded" : "failed");

                    if (success) {
                        Grunt_Survey *gs = (Grunt_Survey *)htLookup(&grunt_survey_map,
                                                                    hashString(resource_id));

                        if (gs != NULL) { // If it has not been already removed
                            Grunt *g = &gt.grunts[htLookup(&gt.map, hashString(node_id))];
                            sbufAdd(gs->grunt_ids, strdup(g->id));
                        }
                    }
                } else if (kafka_message_read->rkt == topic_log) {
                    Char *node_id = cJSON_GetObjectItem(root, "node_id")->valuestring;
                    Char *resource_id = cJSON_GetObjectItem(root, "resource_id")->valuestring;
                    Char *function_id = cJSON_GetObjectItem(root, "function_id")->valuestring;
                    logMessage("Log: %s\n", cJSON_Print(cJSON_Parse((char *)kafka_message_read->payload)));
                    unused_variable(node_id);
                    unused_variable(resource_id);
                    unused_variable(function_id);
                } else {
                    // TODO(naman): Error
                }

                cJSON_Delete(root);
            }

            rd_kafka_message_destroy(kafka_message_read);
        }

        for (Size j = 0; j < gt.map.slot_count; j++) {
            if (gt.map.values[j] != 0) {
                Size index = gt.map.values[j];
                Grunt g = gt.grunts[index];
                if (g.time_to_die <= 0) {
                    logMessage("Deleting grunt: %s\n", g.id);
                    gruntTrackEnd(&gt, g.id);
                }
            }
        }

        for (Size i = 0; i < grunt_survey_map.slot_count; i++) {
            if (grunt_survey_map.keys[i] != 0) {
                Grunt_Survey *gs = (Grunt_Survey *)grunt_survey_map.values[i];

                U64 milli_new = timeMilli();
                gs->milli_passed += milli_new - gs->milli_last;
                gs->milli_last = milli_new;
                if (gs->milli_passed >= conf.grunt_response_wait_time) {
                    Command c = {.kind = Command_RESPONSE_ARBITER_2_DM};
                    c.resource_id = gs->resource_id;

                    for (Size k = 0; k < sbufElemin(gs->grunt_ids); k++) {
                        Size index = htLookup(&gt.map, hashString(gs->grunt_ids[k]));
                        if (index != 0) {
                            Char *id = gt.grunts[index].id;
                            sbufAdd(c.res_a2d.grunt_ids, id);
                        }
                    }

                    sbufAdd(commands, c);

                    sbufDelete(gs->grunt_ids);
                    free(gs->ports);
                    htRemove(&grunt_survey_map, hashString(gs->resource_id));
                }
            }
        }

        for (Size j = 0; j < sbufElemin(commands); j++) {
            Command c = commands[j];

            Sint timestamp = (Sint)time(0);

            Char *output = NULL;
            Char *topic = NULL;

            if (c.kind == Command_REQUEST_ARBITER_2_GRUNT) {
                topic = "REQUEST_RM_2_RD";

                sbufPrint(output, "{\n\"resource_id\": \"%s\"", c.resource_id);
                sbufPrint(output, ",\n\"memory\": %d", c.req_a2g.memory);
                sbufPrint(output, ",\n\"timestamp\": %d", timestamp);
                sbufPrint(output, "\n}\n");
            } else if (c.kind == Command_RESPONSE_ARBITER_2_DM) {
                topic = "RESPONSE_RM_2_DM";

                sbufPrint(output, "{\n\"resource_id\": \"%s\"", c.resource_id);
                sbufPrint(output, ",\n\"timestamp\": %d", timestamp);
                sbufPrint(output, ",\n\"nodes\": [");
                for (Size k = 0; k < sbufElemin(c.res_a2d.grunt_ids); k++) {
                    sbufPrint(output, "\"%s\"", c.res_a2d.grunt_ids[k]);
                    if (k < sbufElemin(c.res_a2d.grunt_ids) - 1) {
                        sbufPrint(output, ",");
                    }
                }
                sbufPrint(output, "]");
                sbufPrint(output, "\n}");
            } else if (c.kind == Command_JOIN_ACK_ARBITER_2_GRUNT) {
                topic = "JOIN_ACK_RM_2_RD";

                sbufPrint(output, "{\n\"node_id\": \"%s\"", c.join_ack_a2g.grunt_id);
                sbufPrint(output, ",\n\"timestamp\": %d", timestamp);
                sbufPrint(output, "\n}");

                free(c.join_ack_a2g.grunt_id);
            } else if (c.kind == Command_REJOIN_ARBITER_2_GRUNT) {
                topic = "REJOIN_RM_2_RD";

                sbufPrint(output, "{\n\"node_id\": \"%s\"", c.rejoin_a2g.grunt_id);
                sbufPrint(output, ",\n\"timestamp\": %d", timestamp);
                sbufPrint(output, "\n}");

                free(c.rejoin_a2g.grunt_id);
            }

            if (output != NULL) {
                printf("Sending to %s\n%s\n", topic, output);
                if (!kafkaWrite(kafka.writer, topic, "resource_manager", output)) {
                    return -1;
                }
            }

            sbufDelete(output);
            free(c.resource_id);

            if (c.kind == Command_RESPONSE_ARBITER_2_DM) {
                for (Size k = 0; k < sbufElemin(c.res_a2d.grunt_ids); k++) {
                    free(c.res_a2d.grunt_ids[k]);
                }
            }
            sbufUnsortedRemove(commands, j);
        }
    }

    for (Size i = 0; i < sbufElemin(kafka.topics); i++) {
        rd_kafka_topic_destroy(kafka.topics[i]);
    }
    rd_kafka_consumer_close(kafka.reader);
    rd_kafka_destroy(kafka.reader);
    for (Size i = 0; i < sbufElemin(kafka.queues); i++) {
        rd_kafka_queue_destroy(kafka.queues[i]);
    }
    rd_kafka_destroy(kafka.writer);

    return 0;
}
