#include <errno.h>
#include <inttypes.h>
#include <rtti.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>

#include "htab.h"
#include "kv.h"
#include "kvmgr.h"
#include "mylock.h"
#include "persist.h"
#include "utils.h"
#include "xce/types.h"
#include "zarray.h"

#include "controller.capnp.h"

typedef struct {
    my_rwlock_t lock;
} xkvmgr_private_t;

typedef struct {
    kvmgr_key_type_t key_type;
    uint32_t         nslots;
    var              kvtab;
    var              data;
} xkvmgr_t;

static size_t xkvmgr_datasize(va_list ap)
{
    size_t sz = sizeof(xkvmgr_t);

    return sz;
}

static int xkvmgr_init(meta_obj_t* obj, va_list ap)
{
    xkvmgr_t*         kvm    = RTTI_OBJ_TO_DATA(xkvmgr_t, obj);
    kvmgr_params_t*   params = va_arg(ap, kvmgr_params_t*);
    kv_params_t       kv_params;
    xzarray_params_t  data_params;
    xkvmgr_private_t* priv = RTTI_OBJ_TO_PRIV(obj);

    if ((params == NULL)
        || ((params->key_type != KVMGR_KEY_INT)
            && (params->key_type != KVMGR_KEY_STRING))
        || (params->nslots <= 1) || (params->delta < 1)) {
        return -EINVAL;
    }

    kv_params.nslots     = params->nslots;
    kv_params.datakiller = NULL;

    switch (params->key_type) {
    case KVMGR_KEY_INT:
        kv_params.keyops = &int_key_ops;
        break;
    case KVMGR_KEY_STRING:
        kv_params.keyops = &string_key_ops;
        break;
    default:
        return -EINVAL;
    }

    kvm->kvtab = makeInstance(&kvtab_type, &kv_params);
    if (kvm->kvtab == NULL) {
        return -errno;
    }

    data_params.delta      = params->delta;
    data_params.datakiller = params->datakiller;
    kvm->data              = makeInstance(&xzarray_type, &data_params);
    if (kvm->data == NULL) {
        destroy(kvm->kvtab);

        return -errno;
    }

    kvm->nslots   = params->nslots;
    kvm->key_type = params->key_type;

    my_rwlock_init(&(priv->lock));

    return 1;
}

static void xkvmgr_destroy(meta_obj_t* obj)
{
    xkvmgr_t*         kvm  = RTTI_OBJ_TO_DATA(xkvmgr_t, obj);
    xkvmgr_private_t* priv = RTTI_OBJ_TO_PRIV(obj);

    if (kvm->data != NULL) {
        destroy(kvm->data);
    }

    if (kvm->kvtab != NULL) {
        destroy(kvm->kvtab);
    }

    my_rwlock_fini(&(priv->lock));
    memset(kvm, 0x0, sizeof(*kvm));

    return;
}

static int xkvmgr_add(meta_obj_t* obj, void* key, void* value)
{
    xkvmgr_t*         kvm = RTTI_OBJ_TO_DATA(xkvmgr_t, obj);
    int               rc;
    xkvmgr_private_t* priv = RTTI_OBJ_TO_PRIV(obj);
    int               idx;

    rc = vstore(kvm->data, value);

    if (rc >= 0) {
        idx = rc;

        rc = key_add(kvm->kvtab, key, (void*)(long)idx);

        if (rc >= 0) {
            rc = idx;
        } else {
            vclear(kvm->data, idx);
        }
    }

    return rc;
}

static int xkvmgr_find(meta_obj_t* obj, void* key, void** value)
{
    xkvmgr_t*         kvm = RTTI_OBJ_TO_DATA(xkvmgr_t, obj);
    int               rc;
    xkvmgr_private_t* priv = RTTI_OBJ_TO_PRIV(obj);
    void*             data;

    *value = NULL;

    rc = key_find(kvm->kvtab, key, &data);
    if (rc >= 0) {
        int idx = (int)(long)data;

        rc = vget(kvm->data, idx, &data);
        if (rc >= 0) {
            *value = data;
            rc     = 1;
        }
    }

    return rc;
}

static int xkvmgr_make_iter(iterator_t* iter, va_list ap)
{
    meta_obj_t*       obj  = iter->obj;
    xkvmgr_t*         kvm  = RTTI_OBJ_TO_DATA(xkvmgr_t, obj);
    xkvmgr_private_t* priv = RTTI_OBJ_TO_PRIV(obj);
    iterator_t*       it   = calloc(1, sizeof(*it));
    int               rc   = 0;

    if (it == NULL) {
        return -ENOMEM;
    }

    rc = iter_make(it, kvm->data);
    if (rc < 0) {
        FREEUP(it);
        return rc;
    }

    iter->data = it;

    return RTTI_ITER_OK;
}

static int xkvmgr_iter_next(iterator_t* iter, void** data)
{
    iterator_t* it = (iterator_t*)iter->data;

    return iter_next(it, data);
}

static int xkvmgr_iter_reset(iterator_t* iter)
{
    iterator_t* it = (iterator_t*)iter->data;

    return iter_reset(it);
}

static void xkvmgr_iter_destroy(iterator_t* iter)
{
    iterator_t*       it   = iter->data;
    meta_obj_t*       obj  = iter->obj;
    xkvmgr_t*         kvm  = RTTI_OBJ_TO_DATA(xkvmgr_t, obj);
    xkvmgr_private_t* priv = RTTI_OBJ_TO_PRIV(obj);

    if (it != NULL) {
        iter_destroy(it);

        FREEUP(it);
    }

    iter->data = NULL;

    return;
}

static int xkvmgr_get_vec(meta_obj_t* obj, int count, int* idxs, void** data)
{
    xkvmgr_t*         kvm  = RTTI_OBJ_TO_DATA(xkvmgr_t, obj);
    xkvmgr_private_t* priv = RTTI_OBJ_TO_PRIV(obj);
    int               rc   = 0;

    rc = vget_vec(kvm->data, count, idxs, data);

    return rc;
}

static int xkvmgr_get(meta_obj_t* obj, int idx, void** data)
{
    xkvmgr_t*         kvm  = RTTI_OBJ_TO_DATA(xkvmgr_t, obj);
    xkvmgr_private_t* priv = RTTI_OBJ_TO_PRIV(obj);
    int               rc   = 0;

    rc = vget(kvm->data, idx, data);

    return rc;
}

static int xkvmgr_put_vec(meta_obj_t* obj, int count, int* idxs, void** data)
{
    xkvmgr_t*         kvm  = RTTI_OBJ_TO_DATA(xkvmgr_t, obj);
    xkvmgr_private_t* priv = RTTI_OBJ_TO_PRIV(obj);
    int               rc   = 0;

    rc = vput_vec(kvm->data, count, idxs, data);

    return rc;
}

static int xkvmgr_put(meta_obj_t* obj, int idx, void* data)
{
    xkvmgr_t*         kvm  = RTTI_OBJ_TO_DATA(xkvmgr_t, obj);
    xkvmgr_private_t* priv = RTTI_OBJ_TO_PRIV(obj);
    int               rc   = 0;

    rc = vput(kvm->data, idx, data);

    return rc;
}

static int xkvmgr_lock(meta_obj_t* obj, uint64_t option)
{
    xkvmgr_private_t* priv = (xkvmgr_private_t*)RTTI_OBJ_TO_PRIV(obj);
    int               rc   = -EINVAL;

    switch (option) {
    case RTTI_LOCK_READ:
        rc = my_rwlock_rlock(&(priv->lock));
        break;
    case RTTI_LOCK_WRITE:
        rc = my_rwlock_wlock(&(priv->lock));
        break;
    default:
        break;
    }

    return rc;
}

static void xkvmgr_unlock(meta_obj_t* obj)
{
    xkvmgr_private_t* priv = (xkvmgr_private_t*)RTTI_OBJ_TO_PRIV(obj);

    my_rwlock_unlock(&(priv->lock));

    return;
}

typeinfo_t xkvmgr_type = {
    .name        = "xkvmgr",
    .privateSize = sizeof(xkvmgr_private_t),
    .dataSize    = (datasize_fn_t)xkvmgr_datasize,
    .init        = (obj_init_fn_t)xkvmgr_init,
    .destroy     = (obj_destroy_fn_t)xkvmgr_destroy,
    .interfaces
    = { .kv     = &(intf_kv_t) { .add_key  = (key_add_fn_t)xkvmgr_add,
                                 .del_key  = (key_del_fn_t)NULL,
                                 .find_key = (key_find_fn_t)xkvmgr_find },
       .varray = &(intf_varray_t) { .vstore   = (vstore_fn_t)NULL,
                                     .vget     = (vget_fn_t)xkvmgr_get,
                                     .vget_vec = (vget_vec_fn_t)xkvmgr_get_vec,
                                     .vput     = (vput_fn_t)xkvmgr_put,
                                     .vput_vec = (vput_vec_fn_t)xkvmgr_put_vec,
                                     .vclear   = (vclear_fn_t)NULL,
                                     .vfreeze  = NULL, /* wired in P0-2 */
                                     .vthaw    = NULL  /* wired in P0-2 */ },
       .iter
        = &(intf_iterator_t) { .make  = (iter_make_fn_t)xkvmgr_make_iter,
                               .next  = (iter_next_fn_t)xkvmgr_iter_next,
                               .reset = (iter_reset_fn_t)xkvmgr_iter_reset,
                               .destroy
                               = (iter_destroy_fn_t)xkvmgr_iter_destroy },
       .lock = &(intf_lock_t) { .lock   = (lock_fn_t)xkvmgr_lock,
                                 .unlock = (unlock_fn_t)xkvmgr_unlock } }
};

/* ------------------------------------------------------------------ */
/* kvmgr_capture_keys — P0-14a                                        */
/* ------------------------------------------------------------------ */

typedef struct {
    slot_key_t*      keys;
    size_t           capacity;
    size_t           filled;
    kvmgr_key_type_t key_type;
    int              oom;
} capture_ctx_t;

static void capture_one(void* raw_key, int index, void* arg)
{
    capture_ctx_t* ctx = (capture_ctx_t*)arg;

    if (ctx->oom || ctx->filled >= ctx->capacity) {
        return;
    }

    slot_key_t* sk = &ctx->keys[ctx->filled];
    sk->index      = index;

    if (ctx->key_type == KVMGR_KEY_STRING) {
        const char* src = (const char*)raw_key;
        strncpy(sk->key, src, XCE_MAXNAME_LEN - 1);
        sk->key[XCE_MAXNAME_LEN - 1] = '\0';
        sk->key_len                   = strlen(sk->key) + 1;
    } else {
        /* Integer key — format as decimal string */
        int64_t intval = (int64_t)(long)raw_key;
        int     n      = snprintf(sk->key, XCE_MAXNAME_LEN,
                                  "%" PRId64, intval);
        if (n < 0 || n >= XCE_MAXNAME_LEN) {
            /* Truncated — mark as oom so caller can detect */
            ctx->oom = 1;
            return;
        }
        sk->key_len = (size_t)n + 1;
    }

    ctx->filled++;
}

/*
 * Count live entries by iterating; kvtab_foreach returns the count.
 * We use a NULL callback to count first, then allocate, then fill.
 */
static void count_one(void* key, int index, void* arg)
{
    (void)key;
    (void)index;
    size_t* n = (size_t*)arg;
    (*n)++;
}

int kvmgr_capture_keys(var kvmgr_var,
                       slot_key_t** keys_out,
                       size_t*      n_keys_out)
{
    *keys_out   = NULL;
    *n_keys_out = 0;

    if (kvmgr_var == NULL) {
        return -EINVAL;
    }

    /* The public API pointer (var) points to the data region;
     * convert to the meta_obj_t header, then extract xkvmgr_t. */
    meta_obj_t* obj = RTTI_DATA_TO_OBJ(kvmgr_var);
    xkvmgr_t*   kvm = RTTI_OBJ_TO_DATA(xkvmgr_t, obj);

    /* Count live entries */
    size_t n = 0;
    kvtab_foreach(kvm->kvtab, count_one, &n);

    if (n == 0) {
        return 0;
    }

    slot_key_t* keys = calloc(n, sizeof(slot_key_t));
    if (keys == NULL) {
        return -ENOMEM;
    }

    capture_ctx_t ctx = {
        .keys     = keys,
        .capacity = n,
        .filled   = 0,
        .key_type = kvm->key_type,
        .oom      = 0,
    };

    kvtab_foreach(kvm->kvtab, capture_one, &ctx);

    if (ctx.oom) {
        free(keys);
        return -ENOMEM;
    }

    *keys_out   = keys;
    *n_keys_out = ctx.filled;
    return 0;
}