#include <arpa/inet.h>
#include <assert.h>
#include <json-c/json.h>
#include <signal.h>
#include <sqlite3.h>
#include <string.h>
#include <threads.h>
#include <zlib.h>

#include "compress.c"
#include "logger.c"
#include "sql.c"

// ----------------------------------------------------------------------------
// Defines
// ----------------------------------------------------------------------------
#define BACKLOG 10
#define EVA_SIZE (sizeof(uint64_t) * MVEC_SIZE)
#define MAX_HM_PIXEL_COUNT 0x400 // must equal DEFVAL_HM_PIXEL_COUNT in client.cpp

// ----------------------------------------------------------------------------
// Declarations
// ----------------------------------------------------------------------------
struct Socket {
    int fd;
    struct sockaddr_in addr;
};

struct CallbackContext {
    struct json_object *response;
    int64_t response_rows;
    int64_t hm_left;
    int64_t hm_pixel_count;
    int64_t hm_pixel_pow;
};

struct RenderContext {
    const struct CallbackContext *callback_context;
    const void *blob;
    size_t blob_size;
    uint64_t eva[EVA_SIZE];
    int64_t out[MAX_HM_PIXEL_COUNT];
};

// ----------------------------------------------------------------------------
// Globals
// ----------------------------------------------------------------------------
struct json_object *g_response_header;
size_t g_blob_count;

// ----------------------------------------------------------------------------
// Event array render function
// ----------------------------------------------------------------------------
int eva_render(void *data) {
    assert(data);

    struct RenderContext *render_context = (struct RenderContext *)data;
    int64_t hm_left = render_context->callback_context->hm_left;
    int64_t hm_pixel_count = render_context->callback_context->hm_pixel_count;
    int64_t hm_pixel_pow = render_context->callback_context->hm_pixel_pow;
    int64_t hm_pixel_res = 1 << hm_pixel_pow;
    const void *blob = render_context->blob;
    size_t blob_size = render_context->blob_size;

#if defined(MVEC_LOOP)
    hm_left %= MVEC_SIZE;
#endif

#if !defined(MVEC_LOOP)
#if !defined(NDEBUG)
    int64_t hm_right = hm_left + hm_pixel_res * hm_pixel_count;
#endif
    assert(hm_left < (int64_t)MVEC_SIZE);
    assert(hm_right <= (int64_t)MVEC_SIZE);
#endif

    // Inflate blob
    struct InflateParams params = {
        .avail_in = blob_size,
        .size = EVA_SIZE,
        .in = (Bytef *)blob,
        .out = (Bytef *)render_context->eva,
    };

    comp_inflate(&params);
    comp_inflate_end(&params);

    for (int64_t i = 0; i < hm_pixel_count; i++) {
        render_context->out[i] = 0l;

        for (int64_t j = 0; j < hm_pixel_res; j++) {
            int64_t coord = hm_left + (i * hm_pixel_res) + j;
#if defined(MVEC_LOOP)
            coord %= MVEC_SIZE;
#endif
            render_context->out[i] += render_context->eva[coord];
        }
    }

    return 0;
}

// ----------------------------------------------------------------------------
// SQL callbacks
// ----------------------------------------------------------------------------
void sql_callback_add_column_name(sqlite3_stmt *sql_stmt, void *data) {
    assert(sql_stmt);
    assert(data);
    assert(sqlite3_column_type(sql_stmt, 0) == SQLITE_TEXT);
    assert(sqlite3_column_type(sql_stmt, 1) == SQLITE_TEXT);
    assert(!strcmp(sqlite3_column_name(sql_stmt, 0), "name"));
    assert(!strcmp(sqlite3_column_name(sql_stmt, 1), "type"));

    const char *col_name = (const char *)sqlite3_column_text(sql_stmt, 0);
    const char *col_type = (const char *)sqlite3_column_text(sql_stmt, 1);
    struct json_object *response_header = (struct json_object *)data;

    if (!json_object_object_get_ex(response_header, col_name, NULL)) {
        json_object_object_add(response_header, col_name, json_object_new_array());
    }

    if (!strcmp(col_type, "BLOB")) {
        g_blob_count++;
    }
}

void sql_callback_add_data(sqlite3_stmt *sql_stmt, void *data) {
    assert(sql_stmt);
    assert(data);

    struct CallbackContext *callback_context = (struct CallbackContext *)data;
    struct json_object *col_data = NULL;
    struct RenderContext *render_contexts = calloc(g_blob_count, sizeof(struct RenderContext));
    thrd_t *threads = calloc(g_blob_count, sizeof(thrd_t));

    for (int i = 0, tid = 0; i < sqlite3_column_count(sql_stmt); i++) {
        if (json_object_object_get_ex(callback_context->response, sqlite3_column_name(sql_stmt, i), &col_data)) {
            if (sqlite3_column_type(sql_stmt, i) == SQLITE_BLOB) {
                render_contexts[tid].callback_context = callback_context;
                render_contexts[tid].blob = sqlite3_value_blob(sqlite3_column_value(sql_stmt, i));
                render_contexts[tid].blob_size = sqlite3_value_bytes(sqlite3_column_value(sql_stmt, i));
                thrd_create(&threads[tid], (thrd_start_t)eva_render, &render_contexts[tid]);
                tid++;
            } else {
                json_object_array_add(col_data, json_object_new_int64(sqlite3_column_int64(sql_stmt, i)));
            }
        }
    }

    for (int i = 0, tid = 0; i < sqlite3_column_count(sql_stmt); i++) {
        if (json_object_object_get_ex(callback_context->response, sqlite3_column_name(sql_stmt, i), &col_data)) {
            if (sqlite3_column_type(sql_stmt, i) == SQLITE_BLOB) {
                thrd_join(threads[tid], NULL);

                for (int64_t j = 0; j < callback_context->hm_pixel_count; j++) {
                    json_object_array_add(col_data, json_object_new_int64(render_contexts[tid].out[j]));
                }

                tid++;
            }
        }
    }

    callback_context->response_rows++;
    log_info("Processed row #%ld", callback_context->response_rows);
    free(render_contexts);
    free(threads);
}

// ----------------------------------------------------------------------------
// Main functions
// ----------------------------------------------------------------------------
void sig_handler(int signo) {
    (void)signo;
    log_warn("Signal received, will stop SALIS data server");
    json_object_put(g_response_header);
    sql_close();
    exit(0);
}

void respond_name(int socket_fd) {
    log_info("Client requested simulation name");

    struct json_object *sim_name = json_object_new_object();
    json_object_object_add(sim_name, "name", json_object_new_string(NAME));
    json_object_to_fd(socket_fd, sim_name, 0);
    json_object_put(sim_name);
}

void respond_opts(int socket_fd) {
    log_info("Client requested simulation options");

    struct json_object *sim_opts = json_object_from_file(SIM_OPTS);
    json_object_to_fd(socket_fd, sim_opts, 0);
    json_object_put(sim_opts);
}

void respond_hash(int socket_fd) {
    log_info("Client requested git hash");

    char buff[41] = { 0 };
    FILE *pipe = popen("git rev-parse HEAD", "r");
    fread(buff, sizeof(char), 40, pipe);
    pclose(pipe);

    struct json_object *git_hash = json_object_new_object();
    json_object_object_add(git_hash, "hash", json_object_new_string(buff));
    json_object_to_fd(socket_fd, git_hash, 0);
    json_object_put(git_hash);
}

void respond_data(int socket_fd, struct json_object *request) {
    assert(request);

    const char *request_str = json_object_to_json_string(request);
    log_info("Client requested simulation data with the following parameters: %s", request_str);
    const char *x_axis = json_object_get_string(json_object_object_get(request, "x-axis"));
    int64_t x_current = json_object_get_int64(json_object_object_get(request, "x-current"));
    int64_t x_high = json_object_get_int64(json_object_object_get(request, "x-high"));
    int64_t nth = json_object_get_int64(json_object_object_get(request, "nth"));
    int64_t entries = json_object_get_int64(json_object_object_get(request, "entries"));

    struct CallbackContext callback_context = {
        .response = NULL,
        .response_rows = 0l,
        .hm_left = json_object_get_int64(json_object_object_get(request, "hm-left")),
        .hm_pixel_count = json_object_get_int64(json_object_object_get(request, "hm-pixel-count")),
        .hm_pixel_pow = json_object_get_int64(json_object_object_get(request, "hm-pixel-pow")),
    };

    json_object_deep_copy(g_response_header, &callback_context.response, NULL);

    const char *x_axis_pref = (!strcmp(x_axis, "rowid") || !strcmp(x_axis, "step")) ? "core." : "";

    sql_exec(
        0, NULL, NULL,
        sql_callback_add_data,
        &callback_context,
        "select * from ("
            "select core.rowid, core.step, * from core inner join arch "
            "where core.rowid = arch.rowid and %s%s > %ld and %s%s <= %ld and core.rowid %% %ld == 0 "
            "order by %s%s desc limit %ld"
        ") order by %s asc;",
        x_axis_pref,
        x_axis,
        x_current,
        x_axis_pref,
        x_axis,
        x_high,
        nth,
        x_axis_pref,
        x_axis,
        entries,
        x_axis
    );

    log_info("Sending client %ld rows of data", callback_context.response_rows);
    json_object_to_fd(socket_fd, callback_context.response, 0);
    json_object_put(callback_context.response);

    shutdown(socket_fd, SHUT_WR);
}

int handle_client(struct Socket *socket) {
    assert(socket);

    char socket_ip[INET_ADDRSTRLEN];
    inet_ntop(AF_INET, &socket->addr.sin_addr, socket_ip, INET_ADDRSTRLEN);
    log_info("Client connected: %s:%d", socket_ip, ntohs(socket->addr.sin_port));

    struct json_object *request_json = json_object_from_fd(socket->fd);
    struct json_object *request_str = NULL;

    if (!json_object_object_get_ex(request_json, "request", &request_str)) assert(false);

    const char *request = json_object_get_string(request_str);
    assert(request);

    if (!strcmp(request, "name")) {
        respond_name(socket->fd);
    } else if (!strcmp(request, "opts")) {
        respond_opts(socket->fd);
    } else if (!strcmp(request, "hash")) {
        respond_hash(socket->fd);
    } else if (!strcmp(request, "data")) {
        respond_data(socket->fd, request_json);
    } else {
        assert(false);
    }

    json_object_put(request_json);

    log_info("Client disconnected: %s:%d", socket_ip, ntohs(socket->addr.sin_port));
    close(socket->fd);

    free(socket);
    return 0;
}

int main(void) {
    log_info("Initializing salis data server");
    log_info("Connecting to database in: %s", DATA_PUSH_PATH);
    sql_open();

    signal(SIGINT, sig_handler);
    signal(SIGTERM, sig_handler);
    signal(SIGPIPE, SIG_IGN); // ignore broken pipes

    log_info("Creating response header");
    g_response_header = json_object_new_object();
    json_object_object_add(g_response_header, "rowid", json_object_new_array());
    sql_exec(
        0, NULL, NULL,
        sql_callback_add_column_name,
        g_response_header,
        "select name, type from pragma_table_info('core') union "
        "select name, type from pragma_table_info('arch');"
    );

    log_info("Found %lu blob columns in database", g_blob_count);

    log_info("Binding to port: %d", PORT);
    int opt = 1;
    int socket_fd = socket(AF_INET, SOCK_STREAM, 0);
    setsockopt(socket_fd, SOL_SOCKET, SO_REUSEADDR, &opt, sizeof(opt));
    struct sockaddr_in socket_addr = { 0 };
    socket_addr.sin_family = AF_INET;
    socket_addr.sin_addr.s_addr = INADDR_ANY;
    socket_addr.sin_port = htons(PORT);
    bind(socket_fd, (struct sockaddr *)&socket_addr, sizeof(struct sockaddr_in));

    listen(socket_fd, BACKLOG);
    log_info("Listening...");

    while (true) {
        struct Socket *socket = calloc(1, sizeof(struct Socket));
        socklen_t socket_len = sizeof(struct sockaddr_in);
        socket->fd = accept(socket_fd, (struct sockaddr *)&socket->addr, &socket_len);

        thrd_t thread;
        thrd_create(&thread, (thrd_start_t)handle_client, socket);
        thrd_detach(thread);
    }

    return 0;
}