#include <assert.h>
#include <errno.h>
#include <rtti.h>
#include <stdatomic.h>
#include <stdlib.h>
#include <string.h>

#include "bitset.h"
#include "mylock.h"
#include "utils.h"
#include "zarray.h"

int zarray_init(zarray_t* array, int delta, zarray_item_killer_t killer)
{
    if ((array == NULL) || (delta <= 0)) {
        return -EINVAL;
    }

    mylock_init(&(array->lock));

    array->delta  = delta;
    array->data   = NULL;
    array->killer = killer;
    array->count  = 0;
    array->caps   = 0;
    array->pos    = 0;
    return 0;
}

static int zarray_ensure_items(zarray_t* array, int mincap)
{
    int   nc, i;
    void* p;

    if (mincap < array->caps) {
        return 1;
    }

    nc = mincap;
    if (nc < array->caps + array->delta) {
        nc = array->caps + array->delta;
    }

    p = (void*)realloc(array->data, nc * sizeof(void*));
    if (p == NULL) {
        return -ENOMEM;
    }

    array->data = p;
    for (i = array->caps; i < nc; i++) {
        array->data[i] = NULL;
    }

    /* this will make we can get next free
     * location quickly
     */
    array->pos  = array->caps;
    array->caps = nc;

    return 1;
}

int zarray_store(zarray_t* array, void* data)
{
    int rc, i;

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

    mylock_lock(&(array->lock));
    rc = zarray_ensure_items(array, array->count + 1);

    if (rc >= 0) {
        int idx;

        rc  = -ENOMEM;
        idx = array->pos;

        for (i = 0; i < array->caps; i++) {
            if (array->data[idx] == NULL) {
                array->data[idx] = data;
                array->count++;
                rc = idx;
                break;
            }

            idx++;
            if (idx >= array->caps) {
                idx = 0;
            }
        }
        if (rc >= 0) {
            array->pos = (rc + 1) % array->caps;
        }
    }

    mylock_unlock(&(array->lock));

    return rc;
}

int zarray_clear(zarray_t* array, int idx)
{
    if (array == NULL) {
        return -EINVAL;
    }

    if ((idx < 0) || (idx >= array->caps)) {
        return -EINVAL;
    }

    mylock_lock(&(array->lock));
    array->data[idx] = NULL;
    mylock_unlock(&(array->lock));

    return 0;
}

void zarray_fini(zarray_t* array)
{
    if (array == NULL) {
        return;
    }

    if (array->data != NULL) {
        int i;

        if (array->killer != NULL) {
            for (i = 0; i < array->caps; i++) {
                if (array->data[i] != NULL) {
                    void* ptr      = array->data[i];
                    array->data[i] = NULL;
                    (array->killer)(ptr);
                }
            }
        }

        free(array->data);
        array->data = NULL;
    }
}

int zarray_iter_ex(zarray_t* array, zarray_iter_skip_fn_t skipfn, xlist_t* l,
                   void* arg)
{
    int i, j;

    if ((array == NULL) || (l == NULL)) {
        return -EINVAL;
    }

    mylock_lock(&(array->lock));

    for (j = 0; j < l->count; j++) {
        zarray_iter_fn_t f = (zarray_iter_fn_t)l->data[j];

        if (f == NULL) {
            continue;
        }

        for (i = 0; i < array->caps; i++) {
            if (skipfn != NULL) {
                if (skipfn(array->data[i])) {
                    continue;
                }
            } else if (array->data[i] == NULL) {
                continue;
            }

            (*f)(arg, array->data[i]);
        }
    }
    mylock_unlock(&(array->lock));

    return 0;
}

int zarray_iter(zarray_t* array, zarray_iter_fn_t f, void* arg)
{
    int     rc;
    xlist_t l;

    if ((array == NULL) || (f == NULL)) {
        return -EINVAL;
    }

    xlist_init(&l, 2);
    xlist_store(&l, f);

    rc = zarray_iter_ex(array, NULL, &l, arg);

    xlist_fini(&l);
    return rc;
}

void* zarray_find(zarray_t* array, zarray_find_cond_fn_t fn, ...)
{
    int     i;
    va_list ap;
    void*   data = NULL;

    if ((array == NULL) || (fn == NULL)) {
        errno = EINVAL;

        return NULL;
    }

    va_start(ap, fn);

    mylock_lock(&(array->lock));
    for (i = 0; i < array->caps; i++) {
        va_list ap1;
        int     rc;

        va_copy(ap1, ap);
        rc = fn(array->data[i], ap1);
        va_end(ap1);

        if (rc) {
            data = array->data[i];
            break;
        }
    }
    mylock_unlock(&(array->lock));
    va_end(ap);

    return data;
}

int zarray_del(zarray_t* array, int idx)
{
    void* ptr = NULL;

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

    if (idx < 0) {
        return -EINVAL;
    }

    if (idx >= array->caps) {
        return 0;
    }

    mylock_lock(&(array->lock));
    ptr              = array->data[idx];
    array->data[idx] = NULL;
    if (ptr != NULL) {
        array->count--;
    }

    /* idx will be a usable position */
    array->pos = idx;
    mylock_unlock(&(array->lock));

    if (array->killer != NULL) {
        if (ptr != NULL) {
            (array->killer)(ptr);
        }
    }

    return 0;
}

int zarray_put_ex(zarray_t* array, int count, int* idxs, void** data)
{
    int i, maxid;
    int rc = 0;

    if ((array == NULL) || (count <= 0) || (idxs == NULL) || (data == NULL)) {
        return -EINVAL;
    }

    maxid = 0;
    for (i = 0; i < count; i++) {
        if (idxs[i] < 0) {
            return -EINVAL;
        }

        if (idxs[i] > maxid) {
            maxid = idxs[i];
        }
    }

    mylock_lock(&(array->lock));
    rc = zarray_ensure_items(array, maxid + 1);
    if (rc >= 0) {
        for (i = 0; i < count; i++) {
            int   idx = idxs[i];
            void* ptr = array->data[idx];

            array->data[idx] = data[i];
            if (ptr != NULL) {
                if (array->killer != NULL) {
                    array->killer(ptr);
                }
            } else {
                array->count++;
            }
        }
    }
    mylock_unlock(&(array->lock));

    return 0;
}

int zarray_put(zarray_t* array, int idx, void* data)
{
    int   idxs[]  = { idx };
    void* datas[] = { data };

    return zarray_put_ex(array, 1, &idxs[0], &datas[0]);
}

int zarray_get_ex(zarray_t* array, int count, int* idxs, void** data)
{
    int i;

    if ((array == NULL) || (count <= 0) || (idxs == NULL) || (data == NULL)) {
        return -EINVAL;
    }

    mylock_lock(&(array->lock));
    for (i = 0; i < count; i++) {
        int idx = idxs[i];

        if (idx >= array->caps) {
            data[i] = NULL;
        } else {
            data[i] = array->data[idx];
        }
    }
    mylock_unlock(&(array->lock));

    return 0;
}

void* zarray_get(zarray_t* array, int idx)
{
    int   idxs[]  = { idx };
    void* datas[] = { NULL };
    int   rc      = 0;

    rc = zarray_get_ex(array, 1, &idxs[0], &datas[0]);
    if (rc < 0) {
        SET_ERROR(-rc);
        return NULL;
    }

    return datas[0];
}

typedef struct {
    _Atomic _Bool inuse;
    void*         data;
} xzarray_item_t;

typedef struct {
    mylock_t        lock;
    xzarray_item_t* shadow;      /* COW snapshot (may be NULL if array was empty at freeze) */
    size_t          shadow_caps; /* capacity of shadow array                                */
    bool            in_cow;      /* true while between freeze and thaw                      */
} xzarray_private_t;

typedef struct {
    int             count;
    size_t          caps;
    size_t          delta;
    int             pos;
    xzarray_item_t* data;
    destroy_fn_t    killer;
} xzarray_t;

static size_t xzarray_datasize(va_list ap)
{
    size_t sz = sizeof(xzarray_t);

    return sz;
}

static int xzarray_init(meta_obj_t* obj, va_list ap)
{
    xzarray_private_t* priv   = (xzarray_private_t*)RTTI_OBJ_TO_PRIV(obj);
    xzarray_t*         a      = RTTI_OBJ_TO_DATA(xzarray_t, obj);
    xzarray_params_t*  params = va_arg(ap, xzarray_params_t*);

    if ((params == NULL) || (params->delta < 1)) {
        return -EINVAL;
    }

    a->count  = 0;
    a->caps   = 0;
    a->delta  = params->delta;
    a->data   = NULL;
    a->pos    = 0;
    a->killer = params->datakiller;

    priv->shadow      = NULL;
    priv->shadow_caps = 0;
    priv->in_cow      = false;

    mylock_init(&(priv->lock));

    return 1;
}

static void xzarray_destroy(meta_obj_t* obj)
{
    xzarray_private_t* priv = (xzarray_private_t*)RTTI_OBJ_TO_PRIV(obj);
    xzarray_t*         a    = RTTI_OBJ_TO_DATA(xzarray_t, obj);
    int                i;

    /* Free shadow if still present (e.g. destroy called mid-COW) */
    if (priv->in_cow) {
        free(priv->shadow);
        priv->shadow      = NULL;
        priv->shadow_caps = 0;
        priv->in_cow      = false;
    }

    for (i = 0; i < (int)a->caps; i++) {
        if (atomic_load_explicit(&a->data[i].inuse, memory_order_acquire)) {
            if (a->killer != NULL) {
                a->killer(a->data[i].data);
            }
        }
    }

    FREEUP(a->data);

    mylock_fini(&(priv->lock));
}

static int xzarray_ensure_items(xzarray_t* array, int mincap)
{
    int   nc, i;
    void* p;

    if (mincap < array->caps) {
        return 1;
    }

    nc = mincap;
    if (nc < array->caps + array->delta) {
        nc = array->caps + array->delta;
    }

    p = (xzarray_item_t*)realloc(array->data, nc * sizeof(xzarray_item_t));
    if (p == NULL) {
        return -ENOMEM;
    }

    array->data = p;
    for (i = array->caps; i < nc; i++) {
        atomic_store_explicit(&array->data[i].inuse, false, memory_order_release);
        array->data[i].data  = NULL;
    }

    /* this will make we can get next free
     * location quickly
     */
    if (atomic_load_explicit(&array->data[array->pos].inuse, memory_order_acquire)) {
        array->pos  = array->caps;
    }

    array->caps = nc;

    return 1;
}

static int xzarray_store(meta_obj_t* obj, void* data)
{
    xzarray_t*         a    = RTTI_OBJ_TO_DATA(xzarray_t, obj);
    xzarray_private_t* priv = (xzarray_private_t*)RTTI_OBJ_TO_PRIV(obj);
    int                rc   = 0;

    rc = xzarray_ensure_items(a, a->count + 1);
    if (rc >= 0) {
        int idx, i;

        rc  = -ENOMEM;
        idx = a->pos;

        for (i = 0; i < a->caps; i++) {
            if (!atomic_load_explicit(&a->data[idx].inuse, memory_order_acquire)) {
                a->data[idx].data = data;
                atomic_store_explicit(&a->data[idx].inuse, true, memory_order_release);
                a->count++;

                rc = idx;
                break;
            }

            idx++;
            if (idx >= a->caps) {
                idx = 0;
            }
        }

        if (rc >= 0) {
            a->pos = (rc + 1) % a->caps;
        }
    }
    return rc;
}

static int xzarray_get_vec(meta_obj_t* obj, int count, int* idxs, void** data)
{
    xzarray_t*         a    = RTTI_OBJ_TO_DATA(xzarray_t, obj);
    xzarray_private_t* priv = (xzarray_private_t*)RTTI_OBJ_TO_PRIV(obj);
    int                i;
    int                rc = 0;

    if ((count <= 0) || (idxs == NULL) || (data == NULL)) {
        return -EINVAL;
    }

    for (i = 0; i < count; i++) {
        int idx = idxs[i];

        /* During COW, check active first, then shadow */
        bool in_active = (idx < (int)a->caps) &&
                         atomic_load_explicit(&a->data[idx].inuse, memory_order_acquire);
        bool in_shadow = !in_active && priv->in_cow &&
                         priv->shadow != NULL &&
                         (idx < (int)priv->shadow_caps) &&
                         atomic_load_explicit(&priv->shadow[idx].inuse, memory_order_acquire);

        if (!in_active && !in_shadow) {
            rc = -ENOENT;
            break;
        }
    }

    if (rc == 0) {
        for (i = 0; i < count; i++) {
            int idx = idxs[i];

            /* Active takes precedence over shadow */
            if ((idx < (int)a->caps) &&
                atomic_load_explicit(&a->data[idx].inuse, memory_order_acquire)) {
                data[i] = a->data[idx].data;
            } else {
                data[i] = priv->shadow[idx].data;
            }
        }
    }
    return rc;
}

static int xzarray_get(meta_obj_t* obj, int idx, void** data)
{
    int idxs[] = { idx };

    return xzarray_get_vec(obj, 1, (int*)&idxs[0], data);
}

static int xzarray_put_vec(meta_obj_t* obj, int count, int* idxs, void** data)
{
    xzarray_t*         a    = RTTI_OBJ_TO_DATA(xzarray_t, obj);
    xzarray_private_t* priv = (xzarray_private_t*)RTTI_OBJ_TO_PRIV(obj);
    int                i, rc, maxid;

    if ((count < 1) || (idxs == NULL) || (data == NULL)) {
        return -EINVAL;
    }

    maxid = 0;
    rc    = 0;

    for (i = 0; i < count; i++) {
        if (idxs[i] < 0) {
            rc = -EINVAL;
            break;
        }
        if (maxid < idxs[i]) {
            maxid = idxs[i];
        }
    }

    if (rc >= 0) {
        rc = xzarray_ensure_items(a, maxid + 1);
    }

    if (rc >= 0) {
        for (i = 0; i < count; i++) {
            int idx = idxs[i];

            if (atomic_load_explicit(&a->data[idx].inuse, memory_order_acquire)) {
                if (a->killer != NULL) {
                    a->killer(a->data[idx].data);
                }
            } else {
                a->count++;
            }

            a->data[idx].data = data[i];
            atomic_store_explicit(&a->data[idx].inuse, true, memory_order_release);
        }
    }

    return rc;
}

static int xzarray_put(meta_obj_t* obj, int idx, void* data)
{
    int   idxs[]  = { idx };
    void* datas[] = { data };
    return xzarray_put_vec(obj, 1, &(idxs[0]), &(datas[0]));
}

static int xzarray_clear(meta_obj_t* obj, int idx)
{
    xzarray_t*         a    = RTTI_OBJ_TO_DATA(xzarray_t, obj);
    xzarray_private_t* priv = (xzarray_private_t*)RTTI_OBJ_TO_PRIV(obj);
    int                rc   = -EINVAL;

    assert(!priv->in_cow && "xzarray_clear called during COW — deferred cleanup required");

    if (idx < (int)a->caps) {
        rc = 0;
        if (atomic_load_explicit(&a->data[idx].inuse, memory_order_acquire)) {
            if (a->killer != NULL) {
                a->killer(a->data[idx].data);
            }
            a->count--;
        }

        atomic_store_explicit(&a->data[idx].inuse, false, memory_order_release);
        a->data[idx].data  = NULL;
        a->pos             = idx; /* we can find the available idx next time */
    }

    return rc;
}

static int xzarray_make_iter(iterator_t* iter, va_list ap)
{
    xzarray_private_t* priv = (xzarray_private_t*)RTTI_OBJ_TO_PRIV(iter->obj);

    iter->data = (void*)(long)0;

    return RTTI_ITER_OK;
}

static int xzarray_iter_next(iterator_t* iter, void** data)
{
    xzarray_t* a = RTTI_OBJ_TO_DATA(xzarray_t, iter->obj);
    int        i;

    for (i = (int)(long)iter->data; i < (int)a->caps; i++) {
        if (atomic_load_explicit(&a->data[i].inuse, memory_order_acquire)) {
            *data      = a->data[i].data;
            iter->data = (void*)(long)(i + 1);

            return RTTI_ITER_OK;
        }
    }

    return RTTI_ITER_NO_DATA;
}

static int xzarray_iter_reset(iterator_t* iter)
{
    iter->data = (void*)(long)0;

    return RTTI_ITER_OK;
}

static int xzarray_lock(meta_obj_t* obj, uint64_t option)
{
    xzarray_private_t* priv = (xzarray_private_t*)RTTI_OBJ_TO_PRIV(obj);

    return mylock_lock(&(priv->lock));
}

static void xzarray_unlock(meta_obj_t* obj)
{
    xzarray_private_t* priv = (xzarray_private_t*)RTTI_OBJ_TO_PRIV(obj);

    mylock_unlock(&(priv->lock));

    return;
}

static void xzarray_iter_destroy(iterator_t* iter)
{
    xzarray_private_t* priv = (xzarray_private_t*)RTTI_OBJ_TO_PRIV(iter->obj);

    return;
}

/* P0-2: COW freeze — O(1) pointer swap, pre-grow active array */
static int xzarray_freeze(meta_obj_t* obj)
{
    xzarray_private_t* priv = (xzarray_private_t*)RTTI_OBJ_TO_PRIV(obj);
    xzarray_t*         a    = RTTI_OBJ_TO_DATA(xzarray_t, obj);

    if (priv->in_cow)
        return -EALREADY; /* already frozen */

    /* Pre-grow active to 1.5x + 16 to avoid realloc during COW (A6-026) */
    size_t          new_caps = a->caps + (a->caps / 2) + 16;
    xzarray_item_t* active   = calloc(new_caps, sizeof(xzarray_item_t));
    if (!active)
        return -ENOMEM;

    priv->shadow      = a->data;   /* may be NULL if array was empty */
    priv->shadow_caps = a->caps;
    priv->in_cow      = true;
    a->data           = active;
    a->caps           = new_caps;
    a->count          = 0;
    /* Start pos past shadow to prevent index collisions with shadow entries */
    a->pos            = (int)priv->shadow_caps;
    return 0;
}

/* P0-2: COW thaw — merge shadow into active, free shadow */
static int xzarray_thaw(meta_obj_t* obj, retain_fn_t should_retain, void* arg)
{
    xzarray_private_t* priv = (xzarray_private_t*)RTTI_OBJ_TO_PRIV(obj);
    xzarray_t*         a    = RTTI_OBJ_TO_DATA(xzarray_t, obj);

    if (!priv->in_cow)
        return -EINVAL; /* not in COW mode */

    /* Merge: for each shadow entry not already written to active, copy it */
    for (size_t i = 0; i < priv->shadow_caps; i++) {
        if (!atomic_load_explicit(&priv->shadow[i].inuse, memory_order_acquire))
            continue;

        void* item = priv->shadow[i].data;

        /* GC: skip entries that should not be retained */
        if (should_retain != NULL && !should_retain(item, arg))
            continue;

        /* Only copy to active if active slot not already updated during COW */
        if (i < a->caps &&
            !atomic_load_explicit(&a->data[i].inuse, memory_order_acquire)) {
            a->data[i].data = item;
            atomic_store_explicit(&a->data[i].inuse, true, memory_order_release);
            a->count++;
        }
        /* If active[i].inuse == true: mutation during COW wins — keep it */
    }

    free(priv->shadow);
    priv->shadow      = NULL;
    priv->shadow_caps = 0;
    priv->in_cow      = false;
    a->pos            = 0;
    return 0;
}

typeinfo_t xzarray_type = {
    .name        = "xzarray",
    .privateSize = sizeof(xzarray_private_t),
    .dataSize    = (datasize_fn_t)xzarray_datasize,
    .init        = (obj_init_fn_t)xzarray_init,
    .destroy     = (obj_destroy_fn_t)xzarray_destroy,
    .interfaces
    = { .varray = &(intf_varray_t) { .vstore   = (vstore_fn_t)xzarray_store,
                                     .vget     = (vget_fn_t)xzarray_get,
                                     .vget_vec = (vget_vec_fn_t)xzarray_get_vec,
                                     .vput     = (vput_fn_t)xzarray_put,
                                     .vput_vec = (vput_vec_fn_t)xzarray_put_vec,
                                     .vclear   = (vclear_fn_t)xzarray_clear,
                                     .vfreeze  = (vfreeze_fn_t)xzarray_freeze,
                                     .vthaw    = (vthaw_fn_t)xzarray_thaw },
       .iter
        = &(intf_iterator_t) { .make  = (iter_make_fn_t)xzarray_make_iter,
                               .next  = (iter_next_fn_t)xzarray_iter_next,
                               .reset = (iter_reset_fn_t)xzarray_iter_reset,
                               .destroy
                               = (iter_destroy_fn_t)xzarray_iter_destroy },
       .lock = &(intf_lock_t) { .lock   = (lock_fn_t)xzarray_lock,
                                 .unlock = (unlock_fn_t)xzarray_unlock } }
};