/* Copyright (c) 2006-2014 Dovecot authors, see the included COPYING file */
/* Copyright (c) 2014 Joshua Atkins <josh@ascendantcom.com> */
/* Copyright (c) 2019-2020 Filip Hanes <filip.hanes@gmail.com> */

#include "lib.h"
#include "array.h"
#include "hash.h"
#include "str.h"
#include "strescape.h"
#include "ioloop.h"
#include "istream.h"
#include "mail-namespace.h"
#include "mail-storage-private.h"
#include "http-url.h"
#include "http-client.h"
#include "settings.h"
#include "fts-elastic-plugin.h"
#include "elastic-connection.h"

#include <json-c/json.h>
#include <stdio.h>


struct elastic_search_context {
    pool_t pool;
    int found;
	struct elastic_result ***results;
    const char *scroll_id;
};


struct elastic_connection {
	struct event *event;
    struct mail_namespace *ns;
    const char *username;

    /* ElasticSearch HTTP API information */
    char *http_host;
    const char *basic_auth_username;
    const char *basic_auth_pass;
    in_port_t http_port;
    char *http_base_path;
    char *http_failure;
    int request_status;

    /* for streaming processing of results */
    struct istream *payload;
    struct io *io;
	struct json_tokener *tok;

    enum elastic_post_type post_type;

    /* context for the current search */
    struct elastic_search_context *ctx;

    /* if we should send ?refresh=true on update _bulk requests */
    unsigned int refresh_on_update:1;
    unsigned int debug:1;
    unsigned int http_ssl:1;
};


int elastic_connection_init(const struct fts_elastic_settings *set,
                            struct mail_namespace *ns,
                            struct elastic_connection **conn_r,
                            const char **error_r,
                            struct event *event_parent)
{
    f_debug("start");
    struct elastic_connection *conn = NULL;
    struct http_url *http_url = NULL;
    const char *error = NULL;

    if (error_r == NULL || set == NULL || conn_r == NULL) {
        i_debug("fts_elastic: error initialising ElasticSearch connection");
        f_debug("return -1");
        return -1;
    }

    /* validate the url */
    if (http_url_parse(set->url, NULL, HTTP_URL_ALLOW_USERINFO_PART,
                pool_datastack_create(), &http_url, &error) < 0) {
        *error_r = t_strdup_printf(
            "fts_elastic: Failed to parse HTTP url: %s", error);
        f_debug("return -1");
        return -1;
    }

    conn = i_new(struct elastic_connection, 1);
	conn->event = event_create(event_parent);
    conn->ctx = i_new(struct elastic_search_context, 1);
    conn->ns = ns;
    conn->username = ns->owner ? ns->owner->username : "-";
    conn->http_host = i_strdup(http_url->host.name);

    if (http_url->user != NULL && http_url->password != NULL) {
        conn->basic_auth_username = i_strdup(http_url->user);
        conn->basic_auth_pass = i_strdup(http_url->password);
    }

    conn->http_port = http_url->port;
    conn->http_base_path = i_strdup(http_url->path);
    conn->http_ssl = http_url->have_ssl;
    conn->debug = set->debug;
    conn->refresh_on_update = set->refresh_on_update;
    conn->tok = json_tokener_new();

    /* guard against init being called multiple times */
    if (elastic_http_client == NULL) {
		settings_event_add_filter_name(conn->event, FTS_ELASTIC_FILTER);
		if (http_client_init_private_auto(conn->event, &elastic_http_client,
						  &error) < 0) {
			*error_r = t_strdup(error);
			return -1;
		}
    }

    *conn_r = conn;

    f_debug("return 0");
    return 0;
}


void elastic_connection_deinit(struct elastic_connection *conn)
{
    f_debug("start");
    if (conn != NULL) {
        i_free(conn->http_host);
        i_free(conn->http_base_path);
        i_free(conn->ctx);
        json_tokener_free(conn->tok);
        event_unref(&conn->event);
        i_free(conn);
    }
    f_debug("end");
}

/* Checks response status code from _bulk request */
static void
elastic_connection_bulk_response(const struct http_response *response,
                                   struct elastic_connection *conn)
{
    f_debug("start");
    if (response != NULL && conn != NULL ) {
        /* 200 OK, 204 continue */
        if (response->status / 100 != 2) {
            i_error("fts_elastic: Indexing failed: %s", response->reason);
            conn->request_status = -1;
        }
    }
    f_debug("end");
}

/* Parses response payload to json from _search request */
static void
elastic_connection_payload_input(struct elastic_connection *conn)
{
    f_debug("start");
    const unsigned char *data = NULL;
    json_object *jobj = NULL;
    enum json_tokener_error jerr;
    size_t size;
    int ret = -1;

    /* continue appending data so long as it is available */
    while ((ret = i_stream_read_data(conn->payload, &data, &size, 0)) > 0) {
        jobj = json_tokener_parse_ex(conn->tok, (const char *)data, size);
        i_stream_skip(conn->payload, size);

        jerr = json_tokener_get_error(conn->tok);
        if (jerr == json_tokener_continue) {
            if (ret < 0)
                i_error("fts_elastic: json response not finished");
        } else if (jerr == json_tokener_success) {
            /* extract values from resulting json object */
            elastic_connection_json(conn, jobj);
        } else {
            i_error("fts_elastic: json tokener error: %s", json_tokener_error_desc(jerr));
            break;
        }
    }

    if (ret == 0) {
        /* we will be called again for more data */
    } else {
        if (conn->payload->stream_errno != 0) {
            i_error("fts_elastic: failed to read payload from HTTP server: %m");
            conn->request_status = -1;
        }

        /* clean-up */
        io_remove(&conn->io);
        i_stream_unref(&conn->payload);
    }
    f_debug("end");
}

/* Parses HTTP response from _search request */
static void
elastic_connection_search_response(const struct http_response *response,
                                   struct elastic_connection *conn)
{
    f_debug("start");
    /* 404's usually mean the index is missing.
     * it could mean you also hit a non-ES service
     * 400 means request json is malformed */
    if (response->status / 100 != 2) {
        i_error("fts_elastic: search failed: %d %s",
                    response->status, response->reason);
        conn->request_status = -1;
    }

    if (response->payload == NULL) {
        i_error("fts_elastic: search failed: empty response payload");
        conn->request_status = -1;
        f_debug("return");
        return;
    }

    i_stream_ref(response->payload);
    conn->payload = response->payload;
    conn->io = io_add_istream(response->payload,
                    elastic_connection_payload_input, conn);
    elastic_connection_payload_input(conn);
    f_debug("end");
}

/* Callback from HTTP request
 * according to post_type decides which callback calls */
static void
elastic_connection_http_response(const struct http_response *response,
                                 struct elastic_connection *conn)
{
    f_debug("start");
    if (response != NULL && conn != NULL) {
        switch (conn->post_type) {
        case ELASTIC_POST_TYPE_SEARCH:
            elastic_connection_search_response(response, conn);
            break;
        case ELASTIC_POST_TYPE_BULK:
            elastic_connection_bulk_response(response, conn);
            break;
        case ELASTIC_POST_TYPE_REFRESH:
        case ELASTIC_POST_TYPE_DELETE:
        case ELASTIC_POST_TYPE_DELETE_BY_QUERY:
            /* not implemented */
            break;
        }
    }
    f_debug("end");
}

/* Performs HTTP POST/DELETE request with callback */
int elastic_connection_post(struct elastic_connection *conn,
                            const char *path, string_t *data)
{
    f_debug("start");
    struct http_client_request *http_req = NULL;
    struct istream *post_payload = NULL;
    const char *method = "POST";

    if (conn == NULL || path == NULL || data == NULL) {
        i_error("fts_elastic: connection_post: critical error during POST");
        f_debug("return -1");
        return -1;
    }

    if (conn->post_type == ELASTIC_POST_TYPE_DELETE) {
        method = "DELETE";
    }

    http_req = http_client_request(elastic_http_client, method, conn->http_host,
                                   path, elastic_connection_http_response, conn);
    http_client_request_set_port(http_req, conn->http_port);
    http_client_request_set_ssl(http_req, conn->http_ssl);
    /* XXX: should be application/x-ndjson for bulk updates, but why when this works? */
    http_client_request_add_header(http_req, "Content-Type", "application/json");
    if (conn->basic_auth_username != NULL && conn->basic_auth_pass != NULL) {
        http_client_request_set_auth_simple(http_req, conn->basic_auth_username, conn->basic_auth_pass);
    }

    post_payload = i_stream_create_from_buffer(data);
    http_client_request_set_payload(http_req, post_payload, TRUE);
    i_stream_unref(&post_payload);
    http_client_request_submit(http_req);

    conn->request_status = 0;
    http_client_wait(elastic_http_client);

    f_debug("return %d", conn->request_status);
    return conn->request_status;
}

/* Iterates over json array hits
 * and fills fts_result->definite_uids
 * and fts->result->scores if present
 */
void elastic_connection_search_hits(struct elastic_search_context *ctx, struct json_object *hits)
{
    f_debug("start");
    struct elastic_result *result;
    ARRAY(struct elastic_result *) results_array;
    struct fts_score_map *scores;
    struct json_object *hit;
    struct json_object *jval;
    uint32_t uid = 0;
    int hits_count = 0;
    int i = 0;
    char *box_guid = "";
    const char *_id;
    char **id_part;
	HASH_TABLE(char *, struct elastic_result *) results_hash;

    if (ctx == NULL || hits == NULL) {
        i_error("fts_elastic: select_json: critical error while processing result JSON");
        f_debug("return");
        return;
    }

    if (json_object_get_type(hits) != json_type_array) {
        i_error("fts_elastic: select_json: response hits are not array");
        f_debug("return");
        return;
    }

	hash_table_create(&results_hash, ctx->pool, 0, str_hash, strcmp);
	p_array_init(&results_array, ctx->pool, 32);

    hits_count = json_object_array_length(hits);
    for (i = 0; i < hits_count; i++) {
        hit = json_object_array_get_idx(hits, i);
        if (!json_object_object_get_ex(hit, "_id", &jval)) {
            i_warning("fts_elastic: key _id not in search response hit:%s",
                        json_object_to_json_string(hit));
            continue;
        }

        _id = json_object_get_string(jval);
        id_part = p_strsplit_spaces(ctx->pool, _id, "/");
        if (str_to_uint32(*id_part, &uid) < 0 || uid == 0) {
            i_warning("fts_elastic: uid <= 0 in _id:\"%s\"", _id);
            continue;
        }
        id_part++;
        if (*id_part == NULL) {
            i_warning("fts_elastic: mailbox guid not found in: %s", _id);
            continue;
        }
        if (strcmp(box_guid, *id_part) != 0) {
            box_guid = *id_part;
            result = hash_table_lookup(results_hash, box_guid);
            if (result == NULL) {
                result = p_new(ctx->pool, struct elastic_result, 1);
                result->box_guid = box_guid;
                p_array_init(&result->uids, ctx->pool, 32);
                p_array_init(&result->scores, ctx->pool, 32);
                hash_table_insert(results_hash, box_guid, result);
                array_push_back(&results_array, &result);
            }
        }
        ctx->found += 1;
        if (seq_range_array_add(&result->uids, uid)) {
            /* duplicate result */
        } else if (json_object_object_get_ex(hit, "_score", &jval)) {
            scores = array_append_space(&result->scores);
            scores->uid = uid;
            scores->score = json_object_get_double(jval);
        }
        /* parse user from _id
        id_part++;
        if (*id_part == NULL) {
            i_warning("fts_elastic: user not found in _id:\"%s\"", _id);
            continue;
        }
        user = p_strdup(ctx->pool, *id_part);
        */
    }
    hash_table_destroy(&results_hash);
    array_append_zero(&results_array);
    *ctx->results = array_front_modifiable(&results_array);
    f_debug("end");
}


/* extract values from resulting json object */
void elastic_connection_json(struct elastic_connection *conn, json_object *jobj)
{
    f_debug("start");
    struct json_object *jvalue = NULL;

    i_assert(jobj != NULL);

    /* Check for error description */
    if (json_object_object_get_ex(jobj, "error", &jvalue)) {
        i_error("fts_elastic: %s", json_object_get_string(jvalue));
        f_debug("return");
        return;
    }

    /* Check if errors are present in response */
    if (json_object_object_get_ex(jobj, "errors", &jvalue)) {
        i_error("fts_elastic: errors in response");
    }

    /* Check if _scroll_id are present in response */
    if (json_object_object_get_ex(jobj, "_scroll_id", &jvalue)) {
        conn->ctx->scroll_id = p_strdup(conn->ctx->pool, json_object_get_string(jvalue));
    }

    switch (conn->post_type) {
    case ELASTIC_POST_TYPE_SEARCH:
        if (!json_object_object_get_ex(jobj, "hits", &jvalue)) {
            i_error("fts_elastic: no .hits in search response");
            break;
        }
        if (!json_object_object_get_ex(jvalue, "hits", &jvalue)) {
            i_error("fts_elastic: no .hits.hits in search response");
            break;
        }
        elastic_connection_search_hits(conn->ctx, jvalue);
        break;
    case ELASTIC_POST_TYPE_BULK:
    case ELASTIC_POST_TYPE_REFRESH:
    case ELASTIC_POST_TYPE_DELETE:
    case ELASTIC_POST_TYPE_DELETE_BY_QUERY:
        /* not implemented */
        break;
    }
    f_debug("end");
}

/* Performs elastic _bulk request
 * checking only response status
 */
int elastic_connection_bulk(struct elastic_connection *conn, string_t *cmd)
{
    f_debug("start");
    const char *path = NULL;

    if (conn == NULL || cmd == NULL) {
        i_error("fts_elastic: connection_bulk: conn or cmd is NULL");
        f_debug("return -1");
        return -1;
    }

    conn->post_type = ELASTIC_POST_TYPE_BULK;
    path = t_strconcat(conn->http_base_path, "_bulk"
                        "?routing=", conn->username,
                        conn->refresh_on_update ? "&refresh=true" : "",
                        NULL);
    elastic_connection_post(conn, path, cmd);
    f_debug("return %d", conn->request_status);
    return conn->request_status;
}


int elastic_connection_refresh(struct elastic_connection *conn)
{
    f_debug("start");
    const char *path = NULL;
    string_t *query = t_str_new_const("", 0);

    if (conn == NULL) {
        i_error("fts_elastic: refresh: critical error");
        f_debug("return -1");
        return -1;
    }

    conn->post_type = ELASTIC_POST_TYPE_REFRESH;
    path = t_strconcat(conn->http_base_path, "_refresh", NULL);
    elastic_connection_post(conn, path, query);

    if (conn->request_status < 0) {
        f_debug("return -1");
        return -1;
    }

    f_debug("return 0");
    return 0;
}

/* Performs elastic search query
 * parses json response
 * and fills fts_result
 */
int elastic_connection_search(struct elastic_connection *conn, pool_t pool, string_t *query, struct elastic_result ***results)
{
    f_debug("start");
    const char *path = NULL;

    if (conn == NULL || query == NULL || results == NULL) {
        i_error("fts_elastic: critical error during search");
        f_debug("return -1");
        return -1;
    }

    i_zero(conn->ctx);
    conn->ctx->pool = pool;
    conn->ctx->results = results;
    conn->ctx->found = 0;
    conn->post_type = ELASTIC_POST_TYPE_SEARCH;

	i_free_and_null(conn->http_failure);
    json_tokener_reset(conn->tok);

    path = p_strconcat(pool, conn->http_base_path, "_search?routing=",
                       conn->username, NULL);
    elastic_connection_post(conn, path, query);

    if (conn->request_status < 0) {
        f_debug("return -1");
        return -1;
    }

    f_debug("return %d", conn->ctx->found);
    return conn->ctx->found;
}

/* Performs elastic search query with scroll
 * parses json response
 * and fills fts_result
 * @deprecated use search_after (optionally with point in time)
 */
int elastic_connection_search_scroll(struct elastic_connection *conn, pool_t pool, string_t *query, struct elastic_result ***results)
{
    f_debug("start");
    static const char *SCROLL_TIMEOUT = "7s";
    const char *path = NULL;

    if (conn == NULL || query == NULL || results == NULL) {
        i_error("fts_elastic: critical error during search scroll");
        f_debug("return -1");
        return -1;
    }

    i_zero(conn->ctx);
    i_assert(conn->ctx != NULL);
    conn->ctx->pool = pool;
    conn->ctx->results = results;
    conn->ctx->found = 0;
    conn->post_type = ELASTIC_POST_TYPE_SEARCH;

	i_free_and_null(conn->http_failure);
    json_tokener_reset(conn->tok);

    path = p_strconcat(pool, conn->http_base_path, "_search?routing=",
                        conn->username, "&scroll=", SCROLL_TIMEOUT, NULL);
    elastic_connection_post(conn, path, query);

    if (conn->ctx->scroll_id == NULL) {
        i_error("fts_elastic: _scroll_id not found in scroll response");
        f_debug("return 0");
        return 0;
    }

    /* TODO: strip index (last segment) from path to allow elastic behind proxy */
    path = "/_search/scroll";
    while (conn->ctx->found >= 9998) {
        conn->ctx->found = 0;
        str_truncate(query, 0);
        str_printfa(query, "{\"scroll\":\"%s\", \"scroll_id\":\"%s\"}",
                            SCROLL_TIMEOUT,
                            conn->ctx->scroll_id);
        elastic_connection_post(conn, path, query);
    }
    /* DELETE search scroll context */
    conn->post_type = ELASTIC_POST_TYPE_DELETE;
    str_truncate(query, 0);
    str_printfa(query, "{\"scroll_id\":\"%s\"}", conn->ctx->scroll_id);
    elastic_connection_post(conn, path, query);

    if (conn->request_status < 0) {
        f_debug("return -1");
        return -1;
    }

    f_debug("return %d", conn->ctx->found);
    return conn->ctx->found;
}

/* Performs elastic search delete by query
 */
int elastic_connection_delete_by_query(struct elastic_connection *conn,
                                       pool_t pool, string_t *query)
{
    f_debug("start");
    const char *path = NULL;

    if (conn == NULL || query == NULL) {
        i_error("fts_elastic: critical error during search scroll");
        f_debug("return -1");
        return -1;
    }

    i_zero(conn->ctx);
    conn->ctx->pool = pool;
    conn->post_type = ELASTIC_POST_TYPE_DELETE_BY_QUERY;

	i_free_and_null(conn->http_failure);

    path = t_strconcat(conn->http_base_path, "_delete_by_query?routing=",
                                                    conn->username, NULL);
    elastic_connection_post(conn, path, query);

    if (conn->request_status < 0) {
        f_debug("return -1");
        return -1;
    }

    f_debug("return %d", conn->ctx->found);
    return conn->ctx->found;
}
