/*
 * Copyright (c) 2017-2021,2024 Red Hat.
 *
 * This library is free software; you can redistribute it and/or modify it
 * under the terms of the GNU Lesser General Public License as published
 * by the Free Software Foundation; either version 2.1 of the License, or
 * (at your option) any later version.
 *
 * This library is distributed in the hope that it will be useful, but
 * WITHOUT ANY WARRANTY; without even the implied warranty of MERCHANTABILITY
 * or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU Lesser General Public
 * License for more details.
 */
#include "schema.h"
#include "batons.h"
#include "slots.h"
#include "util.h"
#include <ctype.h>
#include <search.h>
#ifdef HAVE_STRINGS_H
#include <strings.h>
#endif
#include <valkey/alloc.h>
#include "adlist.h"
#if defined(HAVE_LIBUV)
#include <valkey/adapters/libuv.h>
#else
static int keyClusterLibuvAttach() { return RESP_OK; }
#endif

static char default_server[] = "localhost:6379";

static void
key_server_connect_callback(keysAsyncContext *keys, int status)
{
    if (status == RESP_OK) {
	if (pmDebugOptions.series)
	    fprintf(stderr, "Connected to key server on %s:%d\n",
			keys->c.tcp.host, keys->c.tcp.port);
	keysAsyncEnableKeepAlive((keysAsyncContext *)keys);
	/* TODO: client SSL? inject keysSecureConnection() here */
    } else if (pmDebugOptions.series) {
	if (keys->c.connection_type == RESP_CONN_UNIX)
	    fprintf(stderr, "Connecting to %s failed - %s\n",
			keys->c.unix_sock.path, keys->errstr);
	else
	    fprintf(stderr, "Connecting to %s:%d failed - %s\n",
			keys->c.tcp.host, keys->c.tcp.port, keys->errstr);
    }
}

static void
key_server_disconnect_callback(const keysAsyncContext *keys, int status)
{
    if (status == RESP_OK) {
	if (pmDebugOptions.series)
	    fprintf(stderr, "Disconnected from key server on %s:%d\n",
			keys->c.tcp.host, keys->c.tcp.port);
    } else if (pmDebugOptions.series) {
	if (keys->c.connection_type == RESP_CONN_UNIX)
	    fprintf(stderr, "Disconnecting from %s failed - %s\n",
			keys->c.unix_sock.path, keys->errstr);
	else
	    fprintf(stderr, "Disconnecting from %s:%d failed - %s\n",
			keys->c.tcp.host, keys->c.tcp.port, keys->errstr);
    }
}

void
keySlotsSetupMetrics(keySlots *slots)
{
    pmAtomValue	**table;
    pmUnits	units_count = MMV_UNITS(0, 0, 1, 0, 0, PM_COUNT_ONE);
    pmUnits	units_bytes = MMV_UNITS(1, 0, 0, PM_SPACE_BYTE, 0, 0);
    pmUnits	units_us = MMV_UNITS(0, 1, 0, 0, PM_TIME_USEC, 0);
    void	*map;

    if (slots == NULL || slots->registry == NULL)
	return; /* no metric registry has been set up */

    mmv_stats_add_metric(slots->registry, "requests.total", 1,
	MMV_TYPE_U64, MMV_SEM_COUNTER, units_count, MMV_INDOM_NULL,
	"number of requests",
	"Total number of key server requests sent");

    mmv_stats_add_metric(slots->registry, "requests.error", 2,
	MMV_TYPE_U64, MMV_SEM_COUNTER, units_count, MMV_INDOM_NULL,
	"number of request errors",
	"Total number of key server request errors");

    mmv_stats_add_metric(slots->registry, "responses.total", 3,
	MMV_TYPE_U64, MMV_SEM_COUNTER, units_count, MMV_INDOM_NULL,
	"number of responses",
	"Total number of key server responses received");

    mmv_stats_add_metric(slots->registry, "responses.error", 4,
	MMV_TYPE_U64, MMV_SEM_COUNTER, units_count, MMV_INDOM_NULL,
	"number of error responses",
	"Total number of key server error responses received");

    mmv_stats_add_metric(slots->registry, "responses.time", 5,
	MMV_TYPE_U64, MMV_SEM_COUNTER, units_us, MMV_INDOM_NULL,
	"total time for responses",
	"Cumulative time taken to receive all key server responses");

    mmv_stats_add_metric(slots->registry, "requests.inflight.total", 6,
	MMV_TYPE_U64, MMV_SEM_INSTANT, units_count, MMV_INDOM_NULL,
	"inflight requests",
	"Total number of inflight key server requests");

    mmv_stats_add_metric(slots->registry, "requests.inflight.bytes", 7,
	MMV_TYPE_U64, MMV_SEM_INSTANT, units_bytes, MMV_INDOM_NULL,
	"bytes allocated for inflight requests",
	"Memory currently allocated for inflight key server requests");

    mmv_stats_add_metric(slots->registry, "requests.total_bytes", 8,
	MMV_TYPE_U64, MMV_SEM_COUNTER, units_bytes, MMV_INDOM_NULL,
	"total bytes sent for requests",
	"Cumulative count of bytes sent for all key server requests");

    mmv_stats_add_metric(slots->registry, "responses.total_bytes", 9,
	MMV_TYPE_U64, MMV_SEM_COUNTER, units_bytes, MMV_INDOM_NULL,
	"total bytes received in responses",
	"Cumulative count of bytes received in key server responses");

    slots->map = map = mmv_stats_start(slots->registry);

    table = slots->metrics;
    table[SLOT_REQUESTS_TOTAL] = mmv_lookup_value_desc(map,
					"requests.total", NULL);
    table[SLOT_REQUESTS_ERROR] = mmv_lookup_value_desc(map,
					"requests.error", NULL);
    table[SLOT_RESPONSES_TOTAL] = mmv_lookup_value_desc(map,
					"responses.total", NULL);
    table[SLOT_RESPONSES_ERROR] = mmv_lookup_value_desc(map,
					"responses.error", NULL);
    table[SLOT_RESPONSES_TIME] = mmv_lookup_value_desc(map,
					"responses.time", NULL);
    table[SLOT_REQUESTS_INFLIGHT_TOTAL] = mmv_lookup_value_desc(map,
					"requests.inflight.total", NULL);
    table[SLOT_REQUESTS_INFLIGHT_BYTES] = mmv_lookup_value_desc(map,
					"requests.inflight.bytes", NULL);
    table[SLOT_REQUESTS_TOTAL_BYTES] = mmv_lookup_value_desc(map,
					"requests.total_bytes", NULL);
    table[SLOT_RESPONSES_TOTAL_BYTES] = mmv_lookup_value_desc(map,
					"responses.total_bytes", NULL);
}

int
keySlotsSetMetricRegistry(keySlots *slots, mmv_registry_t *registry)
{
    if (slots) {
	slots->registry = registry;
	return 0;
    }
    return -ENOMEM;
}

keySlotsContext *
keySlotsInit(dict *config, void *events)
{
    keySlotsContext		*context;
    sds			servers = NULL;
    sds			def_servers = NULL;
    sds			username = NULL;
    sds			password = NULL;

    if ((context = (keySlotsContext *)calloc(1, sizeof(keySlotsContext))) == NULL) {
	pmNotifyErr(LOG_ERR, "%s: failed to allocate keySlotsContext\n",
			"keySlotsInit");
	return NULL;
    }

    context->slots.state = SLOTS_DISCONNECTED;
    context->slots.events = events;
    context->slots.keymap = malloc(sizeof(keyMap));
    if (context->slots.keymap == NULL) {
	pmNotifyErr(LOG_ERR, "%s: failed to allocate keymap\n",
		"keySlotsInit");
	free(context);
	return NULL;
    }
    context->slots.keymap->dict = dictCreate(&sdsKeyDictCallBacks);
    context->slots.keymap->privdata = NULL;
    if (context->slots.keymap->dict == NULL) {
	pmNotifyErr(LOG_ERR, "%s: failed to allocate keymap dict\n",
		"keySlotsInit");
	free(context->slots.keymap);
	free(context);
	return NULL;
    }

    servers = pmIniFileLookup(config, "keys", "servers");
    if (servers == NULL)
	servers = pmIniFileLookup(config, "redis", "servers"); // back-compat
    if (servers == NULL)
	servers = pmIniFileLookup(config, "pmseries", "servers"); // back-compat
    if (servers == NULL)
	servers = def_servers = sdsnew(default_server);

    username = pmIniFileLookup(config, "keys", "username");
    if (username == NULL)
	username = pmIniFileLookup(config, "redis", "username"); // back-compat
    if (username == NULL)
	username = pmIniFileLookup(config, "pmseries", "auth.username");
    password = pmIniFileLookup(config, "keys", "password");
    if (password == NULL)
	password = pmIniFileLookup(config, "redis", "password"); // back-compat
    if (password == NULL)
	password = pmIniFileLookup(config, "pmseries", "auth.password");

    /*
     * Add options to the valkeyClusterOptions struct
     * the ini parser already removes spaces at the beginning and end of the
     * configuration values, so checking for empty strings using sdslen() is
     * fine.
     * see note above re empty configuration values having only a password
     * set and no username is a valid key server configuration, details:
     * https://valkey.io/commands/auth
    */

    /* Initialize timeout values in the context structure */
    context->connection_timeout.tv_sec = 5;
    context->connection_timeout.tv_usec = 0;
    context->command_timeout.tv_sec = 60;
    context->command_timeout.tv_usec = 0;

    context->opts.initial_nodes = servers;                    // sds string
    context->opts.username = (username && sdslen(username)) ? username : NULL;
    context->opts.password = (password && sdslen(password)) ? password : NULL;
    context->opts.connect_timeout = &context->connection_timeout;      // struct timeval*
    context->opts.command_timeout = &context->command_timeout;      // struct timeval*
    context->opts.async_connect_callback = key_server_connect_callback;
    context->opts.async_disconnect_callback = key_server_disconnect_callback;
#if defined(HAVE_LIBUV)
    context->opts.attach_fn = keyLibuvAttachAdapter;
    context->opts.attach_data = (uv_loop_t *)events;
#endif

    /* Note: actual connection happens in keySlotsReconnect(), not here.
     * This avoids creating a context that gets immediately freed and recreated. */
    context->slots.acc = NULL;
    
    sdsfree(def_servers); /* Coverity CID370634 */

    return context;
}

/**
 * Handles reconnection logic and schema loading after initial
 * connection (which happens in keySlotsInit with libvalkey)
 */
void
keySlotsReconnect(keySlots *slots, keySlotsFlags flags,
		keysInfoCallBack info, keysDoneCallBack done,
		void *userdata, void *events, void *arg, valkeyClusterOptions *opts)
{
    dictEntry		*entry;
    static int		log_connection_errors = 1;

    if (slots == NULL)
	return;

    slots->state = SLOTS_CONNECTING;
    slots->conn_seq++;

    /* Free old async context if this is a reconnect */
    if (slots->acc != NULL) {
	/* reset key server context in case of reconnect */
	if (slots->acc->err) {
	    /* reset possible 'Connection refused' error before reconnecting */
	    slots->acc->err = 0;
	    memset(slots->acc->errstr, '\0', strlen(slots->acc->errstr));
	}
	keyClusterAsyncDisconnect(slots->acc);

	/* reset keySlots in case of reconnect */
	slots->cluster = 0;
	slots->search = 0;
	/* Clear the keymap by iterating and deleting all entries.
	 * Note: Must free the values (malloc'd int64_t* pointers) before
	 * deleting entries since sdsKeyDictCallBacks has no valDestructor. */
	{
	    dictIterator iter;
	    dictEntry *e;
	    dictInitIterator(&iter, slots->keymap->dict);
	    while ((e = dictNext(&iter)) != NULL) {
		void *val = dictGetVal(e);
		if (val != NULL) {
		    free(val);  /* Free the malloc'd int64_t* position pointer */
		}
		dictDelete(slots->keymap->dict, dictGetKey(e));
	    }
	}

	/* Free old async context completely and create a new one.
	 * libvalkey doesn't support reconnecting an existing async context,
	 * so we need to free the old one and create a new one. */
	keyClusterAsyncFree(slots->acc);
	slots->acc = NULL;
    }
    
    /* Enable blocking initial slot map update for reconnect so the context
     * is immediately ready to accept commands when the function returns. */
    int saved_options = opts->options;
    opts->options |= KEYOPT_BLOCKING_INITIAL_UPDATE;
    
    /* Create new async connection using the same options */
    slots->acc = keyClusterAsyncContextInit(opts);
    
    /* Restore original options */
    opts->options = saved_options;
    
    if (slots->acc == NULL) {
	pmNotifyErr(LOG_ERR, "%s: failed to create new async context\n",
		    "keySlotsReconnect");
	slots->state = SLOTS_DISCONNECTED;
	return;
    }
    if (slots->acc->err == 0) {
	slots->cluster = 1;
    }
    else if (slots->acc->err &&
		strcmp(slots->acc->errstr, RESP_ENOCLUSTER) == 0) {
	/* key server instance has cluster support disabled */
	/* Clear error (following valkeyClusterAsyncClearError pattern) */
	slots->acc->cc.err = 0;
	slots->acc->cc.errstr[0] = '\0';
	slots->acc->err = slots->acc->cc.err;  /* sync acc->err from cc->err */
	slots->cluster = 0;

	/*
	 * Sanity check: show error message if more than one node
	 * is configured, but cluster mode is disabled
	 * otherwise all other nodes silently don't get any data
	 */
	{
	    dictIterator iter;
	    dictInitIterator(&iter, slots->acc->cc.nodes);
	    entry = dictNext(&iter);
	    if (entry && dictNext(&iter)) {
		pmNotifyErr(LOG_ERR, "%s: more than one node is configured, "
			    "but cluster mode is disabled", "keySlotsReconnect");
		slots->state = SLOTS_ERR_FATAL;
		return;
	    }
	}
    }
    else {
	if (log_connection_errors || pmDebugOptions.desperate) {
	    /* acc->errstr should point to acc->cc.errstr, but be defensive */
	    const char *errstr = NULL;
	    if (slots->acc->errstr != NULL && slots->acc->errstr[0] != '\0') {
		errstr = slots->acc->errstr;
	    } else if (slots->acc->cc.errstr[0] != '\0') {
		errstr = slots->acc->cc.errstr;
	    } else {
		/* Check if we can get more info from errno or error code */
		switch (slots->acc->err) {
		    case RESP_ERR_IO:
			errstr = "Connection refused";
			break;
		    default:
			errstr = "Unknown error";
			break;
		}
	    }
	    pmNotifyErr(LOG_INFO, "Cannot connect to key server: %s\n", errstr);
	    log_connection_errors = 0;
	}
	slots->state = SLOTS_DISCONNECTED;
	return;
    }

    slots->state = SLOTS_CONNECTED;
    log_connection_errors = 1;
    keysSchemaLoad(slots, flags, info, done, userdata, events, arg);
}

/**
 * this method allocates the keySlots struct and exists for backwards
 * compatibility, the actual connection to the key server happens in
 * keySlotsReconnect()
 */
keySlotsContext *
keySlotsConnect(dict *config, keySlotsFlags flags,
		keysInfoCallBack info, keysDoneCallBack done,
		void *userdata, void *events, void *arg)
{
    keySlotsContext		*context;
    sds				enabled, msg;

    if (!(enabled = pmIniFileLookup(config, "resp", "enabled")))
	enabled = pmIniFileLookup(config, "redis", "enabled"); // back-compat
    if (enabled && strcmp(enabled, "false") == 0)
	return NULL;

    context = keySlotsInit(config, events);
    if (context == NULL) {
	msg = NULL;
	infofmt(msg, "Failed to allocate memory for key server slots");
	info(PMLOG_ERROR, msg, arg);
	sdsfree(msg);
	return NULL;
    }

    keySlotsReconnect(&context->slots, flags, info, done, userdata, events, arg, &context->opts);
    return context;
}

void
keySlotsFree(keySlots *slots)
{
    keyClusterAsyncDisconnect(slots->acc);
    keyClusterAsyncFree(slots->acc);
    if (slots->keymap) {
	if (slots->keymap->privdata)
	    sdsfree((sds)slots->keymap->privdata);
	/* Free all values (malloc'd int64_t* pointers) before releasing dict.
	 * sdsKeyDictCallBacks has no valDestructor, so we must do this manually. */
	if (slots->keymap->dict) {
	    dictIterator iter;
	    dictEntry *e;
	    dictInitIterator(&iter, slots->keymap->dict);
	    while ((e = dictNext(&iter)) != NULL) {
		void *val = dictGetVal(e);
		if (val != NULL) {
		    free(val);
		}
	    }
	    dictRelease(slots->keymap->dict);
	}
	free(slots->keymap);
    }
    memset(slots, 0, sizeof(*slots));
    free(slots);
}

void
keySlotsContextFree(keySlotsContext *context)
{
    if (context == NULL)
	return;

    /* Set state to disconnected BEFORE invoking callbacks to prevent
     * callbacks from issuing new requests during shutdown */
    context->slots.state = SLOTS_DISCONNECTED;

    keyClusterAsyncDisconnect(context->slots.acc);
    keyClusterAsyncFree(context->slots.acc);
    if (context->slots.keymap) {
	if (context->slots.keymap->privdata)
	    sdsfree((sds)context->slots.keymap->privdata);
	/* Free all values (malloc'd int64_t* pointers) before releasing dict.
	 * sdsKeyDictCallBacks has no valDestructor, so we must do this manually. */
	if (context->slots.keymap->dict) {
	    dictIterator iter;
	    dictEntry *e;
	    dictInitIterator(&iter, context->slots.keymap->dict);
	    while ((e = dictNext(&iter)) != NULL) {
		void *val = dictGetVal(e);
		if (val != NULL) {
		    free(val);
		}
	    }
	    dictRelease(context->slots.keymap->dict);
	}
	free(context->slots.keymap);
    }
    memset(context, 0, sizeof(*context));
    free(context);
}

static inline uint64_t
gettimeusec(void)
{
    struct timeval now;
    if (gettimeofday(&now, NULL) < 0)
        return -1;
    return (uint64_t)now.tv_sec * 1000000 + (uint64_t)now.tv_usec;
}

static keySlotsReplyData *
keySlotsReplyDataAlloc(keySlots *slots, size_t req_size,
			keyClusterCallbackFn *callback, void *arg)
{
    keySlotsReplyData *srd;

    srd = calloc(1, sizeof(keySlotsReplyData));
    if (srd == NULL) {
        return NULL;
    }

    srd->slots = slots;
    srd->conn_seq = slots->conn_seq;
    srd->start = gettimeusec();
    srd->req_size = req_size;
    srd->callback = callback;
    srd->arg = arg;
    return srd;
}

static inline void
keySlotsReplyDataFree(keySlotsReplyData *srd)
{
    free(srd);
}

uint64_t
keySlotsInflightRequests(keySlots *slots)
{
    pmAtomValue		*atom;

    atom = slots ? slots->metrics[SLOT_REQUESTS_INFLIGHT_TOTAL] : NULL;
    return atom ? atom->ull : 0;
}

static void
keySlotsReplyCallback(keyClusterAsyncContext *c, void *r, void *arg)
{
    keySlotsReplyData *srd = arg;
    respReply 		*reply = r;
    void		*map = srd->slots->map;

    if (map) {
	pmAtomValue	**metrics = srd->slots->metrics;
	pmAtomValue	*value;
	uint64_t	delta = gettimeusec();

	delta = (delta < srd->start) ? 0 : delta - srd->start;
	mmv_add(map, metrics[SLOT_RESPONSES_TIME], &delta);

	value = metrics[SLOT_REQUESTS_INFLIGHT_TOTAL];
	delta = value && value->ull > 0 ? value->ull - 1 : 0;
	mmv_set(map, metrics[SLOT_REQUESTS_INFLIGHT_TOTAL], &delta);

	value = metrics[SLOT_REQUESTS_INFLIGHT_BYTES];
	delta = value ? value->ull : 0;
	delta = (delta < srd->req_size) ? 0 : delta - srd->req_size;
	mmv_set(map, metrics[SLOT_REQUESTS_INFLIGHT_BYTES], &delta);

	delta = srd->req_size;
	mmv_add(map, metrics[SLOT_RESPONSES_TOTAL_BYTES], &delta);
	mmv_inc(map, metrics[SLOT_RESPONSES_TOTAL]);

	if (reply == NULL || reply->type == RESP_REPLY_ERROR)
	    mmv_inc(map, metrics[SLOT_RESPONSES_ERROR]);
    }

    /**
     * handle connection resets
     *
     * Why here and not in key_server_disconnect_callback?
     * Access to keySlots->state is required. We cannot save a pointer to
     * keySlots in keysAsyncContext->data, because this member is already
     * used by hiredis-cluster
     *
     * fwiw, in case one node of a cluster is down, slots->state should not be
     * set to SLOTS_DISCONNECTED, because there (should) be a failover in place
     * and another node will handle the requests.
     *
     * A future improvement would be to update hiredis-cluster and add a 'data'
     * member to the async cluster context, analogue to the 'data' member of
     * hiredis, update the disconnect callback to return the cluster context and
     * use this new disconnect callback instead of the conditional below.
     *
     * Register a server disconnect if:
     * * The server returns an I/O error.  In this case errno is also set, but
     *   there are lots of different error codes for connection failures (for example
     *   ECONNRESET, ENETUNREACH, ENETDOWN, ...) - defensively assume all require a
     *   reconnect
     * * Server returns the "LOADING ... is loading the dataset in memory" error
     * * Ignore any errors for server requests pre-dating the latest (current)
     *   connection (to handle the case where a callback returns after a new
     *   connection was already established)
     * * Ignore any errors if the state is already set to SLOTS_DISCONNECTED
     * * Ignore errors if cluster mode is enabled.
     */
    if (((reply == NULL && c->err == RESP_ERR_IO) ||
         (reply != NULL && reply->type == RESP_REPLY_ERROR &&
	  (strncmp(reply->str, RESP_ELOADING, strlen(RESP_ELOADING)) == 0 &&
	   strstr(reply->str, RESP_ELOADDATA) != NULL))) &&
	srd->conn_seq == srd->slots->conn_seq &&
	srd->slots->state != SLOTS_DISCONNECTED &&
	srd->slots->cluster == 0) {
	pmNotifyErr(LOG_ERR, "Lost connection to key server.\n");
	srd->slots->state = SLOTS_DISCONNECTED;
    }

    srd->callback(c, r, srd->arg);
    keySlotsReplyDataFree(arg);
}

/*
 * Submit an arbitrary request to a (set of) key server instance(s).
 * The given key is used to determine the slot used, as per the
 * cluster specification - https://valkey.io/topics/cluster-spec
 * 
 * Serves mainly as a wrapper to keyClusterAsyncFormattedCommand
 * including debug output and error handling
 */
int
keySlotsRequest(keySlots *slots, const sds cmd,
		keyClusterCallbackFn *callback, void *arg)
{
    int			sts;
    uint64_t		size;
    keySlotsReplyData	*srd;

    /*
     * keySlotsSetupStart() also sends key server requests,
     * therefore both SLOTS_CONNECTED and SLOTS_READY states are valid
     */
    if (UNLIKELY(slots->state != SLOTS_CONNECTED && slots->state != SLOTS_READY))
	return -ENOTCONN;

    if (!slots->cluster)
	return keySlotsRequestFirstNode(slots, cmd, callback, arg);

    if (UNLIKELY(pmDebugOptions.desperate))
	fprintf(stderr, "%s: sending raw key server command:\n%s",
			"keySlotsRequest", cmd);

    size = sdslen(cmd);
    if ((srd = keySlotsReplyDataAlloc(slots, size, callback, arg)) == NULL) {
	mmv_inc(slots->map, slots->metrics[SLOT_REQUESTS_ERROR]);
	pmNotifyErr(LOG_ERR, "%s: failed to allocate reply data (%llu bytes)\n",
			"keySlotsRequest", (unsigned long long)size);
	return -ENOMEM;
    }
    if ((sts = keyClusterAsyncFormattedCommand(slots->acc,
		    keySlotsReplyCallback, srd, cmd, size)) != RESP_OK) {
	mmv_inc(slots->map, slots->metrics[SLOT_REQUESTS_ERROR]);
	pmNotifyErr(LOG_ERR, "%s: %s (%s)\n", "keySlotsRequest",
		slots->acc->errstr, cmd);
	keySlotsReplyDataFree(srd);
	return -ENOMEM;
    }

    mmv_add(slots->map, slots->metrics[SLOT_REQUESTS_INFLIGHT_BYTES], &size);
    mmv_add(slots->map, slots->metrics[SLOT_REQUESTS_TOTAL_BYTES], &size);
    mmv_inc(slots->map, slots->metrics[SLOT_REQUESTS_INFLIGHT_TOTAL]);
    mmv_inc(slots->map, slots->metrics[SLOT_REQUESTS_TOTAL]);

    return RESP_OK;
}

int
keySlotsRequestFirstNode(keySlots *slots, const sds cmd,
		keyClusterCallbackFn *callback, void *arg)
{
    dictEntry		*entry;
    valkeyClusterNode	*node;
    keySlotsReplyData	*srd;
    uint64_t		size;
    int			sts;

    /*
     * keySlotsSetupStart() also sends key server requests,
     * therefore both SLOTS_CONNECTED and SLOTS_READY states are valid
     */
    if (UNLIKELY(slots->state != SLOTS_CONNECTED && slots->state != SLOTS_READY))
	return -ENOTCONN;

    {
	dictIterator iter;
	dictInitIterator(&iter, slots->acc->cc.nodes);
	entry = dictNext(&iter);
    }
    if (!entry) {
	pmNotifyErr(LOG_ERR, "%s: No key server node configured.",
			"keySlotsRequestFirstNode");
	return RESP_ERR;
    }

    node = dictGetVal(entry);
    if (UNLIKELY(pmDebugOptions.desperate))
	fprintf(stderr, "%s: sending raw key server command to node %s\n%s",
			"keySlotsRequestFirstNode", node->addr, cmd);

    size = sdslen(cmd);
    if ((srd = keySlotsReplyDataAlloc(slots, size, callback, arg)) == NULL) {
	mmv_inc(slots->map, slots->metrics[SLOT_REQUESTS_ERROR]);
	pmNotifyErr(LOG_ERR, "%s: failed to allocate reply data (%llu bytes)",
			"keySlotsRequestFirstNode", (unsigned long long)size);
	return -ENOMEM;
    }
    sts = keyClusterAsyncFormattedCommandToNode(slots->acc, node,
		keySlotsReplyCallback, srd, cmd, size);
    if (sts != RESP_OK) {
	mmv_inc(slots->map, slots->metrics[SLOT_REQUESTS_ERROR]);
	pmNotifyErr(LOG_ERR, "%s: %s (%s)\n",
		"keySlotsRequestFirstNode", slots->acc->errstr, cmd);
	keySlotsReplyDataFree(srd);
	return -ENOMEM;
    }

    mmv_add(slots->map, slots->metrics[SLOT_REQUESTS_INFLIGHT_BYTES], &size);
    mmv_add(slots->map, slots->metrics[SLOT_REQUESTS_TOTAL_BYTES], &size);
    mmv_inc(slots->map, slots->metrics[SLOT_REQUESTS_INFLIGHT_TOTAL]);
    mmv_inc(slots->map, slots->metrics[SLOT_REQUESTS_TOTAL]);

    return RESP_OK;
}

int
keySlotsProxyConnect(keySlots *slots, keysInfoCallBack info,
	respReader **readerp, const char *buffer, ssize_t nread,
	keyClusterCallbackFn *callback, void *arg)
{
    respReader		*reader = *readerp;
    respReply		*reply = NULL;
    dictEntry		*entry;
    size_t		replyStartPosition;
    long long		position;
    sds			cmd, msg = NULL;
    int			hasKey;
    int			sts;

    if (!reader &&
	(reader = *readerp = respReaderCreate()) == NULL) {
	infofmt(msg, "out-of-memory for key client reader");
	info(PMLOG_REQUEST, msg, arg), sdsfree(msg);
	return -ENOMEM;
    }

    if (respReaderFeed(reader, buffer, nread) != RESP_OK) {
	infofmt(msg, "failed to parse RESP protocol request");
	info(PMLOG_REQUEST, msg, arg), sdsfree(msg);
	return -EPROTO;
    }

    /* parse all key server requests contained in buffer (pipelining) */
    while (1) {
	replyStartPosition = reader->pos;
	sts = respReaderGetReply(reader, (void **)&reply);
	if (sts != RESP_OK) {
	    infofmt(msg, "failed to parse RESP protocol request");
	    info(PMLOG_REQUEST, msg, arg), sdsfree(msg);
	    return -EPROTO;
	}
	if (reply == NULL) {
	    break;
	}

	cmd = NULL;
	hasKey = 0;
	if (reply->type == RESP_REPLY_ARRAY ||
	    reply->type == RESP_REPLY_MAP ||
	    reply->type == RESP_REPLY_SET)
	    cmd = sdsnew(reply->element[0]->str);
	if (cmd && (entry = dictFind(slots->keymap->dict, cmd)) != NULL) {
	    position = *(int64_t *)dictGetVal(entry);
	    if (position > 0 && position < reply->elements)
		hasKey = 1;
	}
	sdsfree(cmd);

	cmd = sdsnewlen(reader->buf + replyStartPosition, reader->pos - replyStartPosition);
	if (hasKey)
	    sts = keySlotsRequest(slots, cmd, callback, arg);
	else
	    sts = keySlotsRequestFirstNode(slots, cmd, callback, arg);
	sdsfree(cmd);

	if (sts != RESP_OK) {
	    respReply *errorReply = calloc(1, sizeof(respReply));
	    errorReply->type = RESP_REPLY_ERROR;
	    errorReply->str = slots->acc->errstr;
	    errorReply->len = strlen(slots->acc->errstr);
	    callback(slots->acc, errorReply, arg);
	}
    }
    return 0;
}

void
keySlotsProxyFree(respReader *reader)
{
    if (reader)
	respReaderFree(reader);
}

/*
 * Helper routines for handling various expected RESP reply types.
 */

int
testReplyError(respReply *reply, const char *server_message)
{
    return (reply && reply->type == RESP_REPLY_ERROR &&
	    strcmp(reply->str, server_message) == 0);
}

void
reportReplyError(keysInfoCallBack info, void *userdata,
	keyClusterAsyncContext *acc, respReply *reply, const char *format, va_list argp)
{
    sds			msg;

    msg = sdscatvprintf(sdsempty(), format, argp);
    if (reply && reply->type == RESP_REPLY_ERROR)
	msg = sdscatfmt(msg, "\nRESP reply error: %s", reply->str);
    else if (acc->err)
	msg = sdscatfmt(msg, "\nRESP acc error: %s", acc->errstr);
    else if (acc->cc.err)
	msg = sdscatfmt(msg, "\nRESP cc error: %s", acc->cc.errstr);
    info(PMLOG_RESPONSE, msg, userdata);
    sdsfree(msg);
}

int
checkStatusReplyOK(keysInfoCallBack info, void *userdata,
		keyClusterAsyncContext *acc, respReply *reply, const char *format, ...)
{
    va_list		argp;

    if (reply && reply->type == RESP_REPLY_STATUS &&
	(strcmp("OK", reply->str) == 0 || strcmp("QUEUED", reply->str) == 0))
	return 0;
    va_start(argp, format);
    reportReplyError(info, userdata, acc, reply, format, argp);
    va_end(argp);
    return -1;
}

int
checkStreamReplyString(keysInfoCallBack info, void *userdata,
	keyClusterAsyncContext *acc, respReply *reply, sds s, const char *format, ...)
{
    va_list		argp;

    if (reply && reply->type == RESP_REPLY_STRING && strcmp(s, reply->str) == 0)
	return 0;
    va_start(argp, format);
    reportReplyError(info, userdata, acc, reply, format, argp);
    va_end(argp);
    return -1;
}

int
checkArrayReply(keysInfoCallBack info, void *userdata,
	keyClusterAsyncContext *acc, respReply *reply, const char *format, ...)
{
    va_list		argp;

    if (reply && reply->type == RESP_REPLY_ARRAY)
	return 0;
    va_start(argp, format);
    reportReplyError(info, userdata, acc, reply, format, argp);
    va_end(argp);
    return -1;
}

long long
checkIntegerReply(keysInfoCallBack info, void *userdata,
	keyClusterAsyncContext *acc, respReply *reply, const char *format, ...)
{
    va_list		argp;

    if (reply && reply->type == RESP_REPLY_INTEGER)
	return reply->integer;
    va_start(argp, format);
    reportReplyError(info, userdata, acc, reply, format, argp);
    va_end(argp);
    return -1;
}

sds
checkStringReply(keysInfoCallBack info, void *userdata,
	keyClusterAsyncContext *acc, respReply *reply, const char *format, ...)
{
    va_list		argp;

    if (reply && reply->type == RESP_REPLY_STRING)
	return sdsnew(reply->str);
    va_start(argp, format);
    reportReplyError(info, userdata, acc, reply, format, argp);
    va_end(argp);
    return NULL;
}
