#include <errno.h>
#include <rtti.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/types.h>

#include "bitset.h"
#include "utils.h"

#if defined(__GNUC__)
#define FFSL(x) __builtin_ffsl(x)
#else
#define FFSL(x) ffsl(x)
#endif

typedef unsigned long bitset_page_data_t;

#define BITSET_PAGE_SIZE (4096)

#define BITSET_ITEM_BITS (sizeof(bitset_page_data_t) * 8)

#define BITSET_PAGE_ITEMS (BITSET_PAGE_SIZE / sizeof(bitset_page_data_t))

#define BITSET_PAGE_BITS (BITSET_PAGE_SIZE * 8)

typedef struct {
    bitset_page_data_t data[BITSET_PAGE_ITEMS];
} bitset_page_t;

#define BIT_PAGE(idx)   ((idx) / BITSET_PAGE_BITS)
#define BIT_OFFSET(idx) (((idx) % BITSET_PAGE_BITS) / BITSET_ITEM_BITS)
#define BIT_BIT(idx)    ((idx) % BITSET_ITEM_BITS)

#define BIT_IDX(pg, offset, bit)                                               \
    ((pg) * BITSET_PAGE_BITS + (offset) * BITSET_ITEM_BITS + (bit))

typedef struct {
    int             npages;
    bitset_page_t** pages;
} bitset_t;

static size_t bitset_datasize(va_list ap)
{
    size_t size = sizeof(bitset_t);

    return size;
}

static int bitset_init(meta_obj_t* obj, va_list ap)
{
    int       max_pages = va_arg(ap, int);
    bitset_t* bitset    = RTTI_OBJ_TO_DATA(bitset_t, obj);

    if (max_pages <= 0) {
        return -EINVAL;
    }

    bitset->pages = (bitset_page_t**)calloc(max_pages, sizeof(bitset_page_t*));
    if (bitset->pages == NULL) {
        return -ENOMEM;
    }

    bitset->npages = max_pages;

    return 1;
}

static void bitset_destroy(meta_obj_t* obj)
{
    bitset_t* bitset = RTTI_OBJ_TO_DATA(bitset_t, obj);
    int       i;

    if (bitset->pages == NULL) {
        return;
    }

    for (i = 0; i < bitset->npages; i++) {
        FREEUP(bitset->pages[i]);
    }

    FREEUP(bitset->pages);
}

static int bitset_test(meta_obj_t* obj, unsigned long idx)
{
    bitset_t*      bitset = RTTI_OBJ_TO_DATA(bitset_t, obj);
    int            pg     = BIT_PAGE(idx);
    int            offset = BIT_OFFSET(idx);
    int            bit    = BIT_BIT(idx);
    bitset_page_t* page   = NULL;

    if (pg >= bitset->npages) {
        return -EINVAL;
    }

    page = bitset->pages[pg];
    if (page == NULL) {
        return 0;
    }

    if (page->data[offset] == 0) {
        return 0;
    }

    return ((page->data[offset] & (((bitset_page_data_t)0x01) << bit)) != 0);
}

static int bitset_set(meta_obj_t* obj, unsigned long idx)
{
    bitset_t*      bitset = RTTI_OBJ_TO_DATA(bitset_t, obj);
    int            pg     = BIT_PAGE(idx);
    int            offset = BIT_OFFSET(idx);
    int            bit    = BIT_BIT(idx);
    bitset_page_t* page   = NULL;

    if (pg >= bitset->npages) {
        return -EINVAL;
    }

    page = bitset->pages[pg];
    if (page == NULL) {
        page = (bitset_page_t*)calloc(1, sizeof(*page));
        if (page == NULL) {
            return -ENOMEM;
        }

        bitset->pages[pg] = page;
    }

    page->data[offset]
        = page->data[offset] | (((bitset_page_data_t)0x01) << bit);

    return 1;
}

static int bitset_clr(meta_obj_t* obj, unsigned long idx)
{
    bitset_t*      bitset = RTTI_OBJ_TO_DATA(bitset_t, obj);
    int            pg     = BIT_PAGE(idx);
    int            offset = BIT_OFFSET(idx);
    int            bit    = BIT_BIT(idx);
    bitset_page_t* page   = NULL;

    if (pg >= bitset->npages) {
        return -EINVAL;
    }

    page = bitset->pages[pg];
    if (page == NULL) {
        return 0;
    }

    page->data[offset]
        = page->data[offset] & ~(((bitset_page_data_t)0x01) << bit);

    return 1;
}

static int bitset_clrall(meta_obj_t* obj)
{
    bitset_t* bitset = RTTI_OBJ_TO_DATA(bitset_t, obj);
    int       pg     = 0;

    for (pg = 0; pg < bitset->npages; pg++) {
        bitset_page_t* page = bitset->pages[pg];

        if (page == NULL) {
            continue;
        }

        memset(page, 0x0, sizeof(*page));
    }
    return 0;
}

typedef struct {
    int overbit;
    int pg;
    int offset;
    int bit;
} iterator_data_t;

static int bitset_iter_make(iterator_t* iter, va_list ap)
{
    int              overbit = va_arg(ap, int);
    iterator_data_t* p       = NULL;

    if ((overbit != BITSET_ITER_OVER_0) && (overbit != BITSET_ITER_OVER_1)) {
        return -EINVAL;
    }

    p = (iterator_data_t*)calloc(1, sizeof(*p));
    if (p == NULL) {
        return -ENOMEM;
    }

    p->overbit = overbit;
    p->pg      = 0;
    p->offset  = 0;
    p->bit     = 0;

    iter->data = (void*)p;

    return 0;
}

static unsigned long bitset_next_over0(bitset_t* bitset, int pg, int offset,
                                       int bit)
{
    while (pg < bitset->npages) {
        while (offset < BITSET_PAGE_ITEMS) {
            bitset_page_data_t val = ~(bitset->pages[pg]->data[offset] >> bit);
            bitset_page_data_t mask = ((~(bitset_page_data_t)0) >> bit);

            val = val & mask;
            
            if (val != 0) {
                int bit1 = FFSL(val);

                if (bit1 != 0) {
                    return BIT_IDX(pg, offset, bit + bit1 - 1);
                }
            }

            offset++;
            bit = 0;
        }

        pg++;
        offset = 0;
        bit    = 0;
    }

    return BIT_IDX(bitset->npages, 0, 0);
}

static unsigned long bitset_next_over1(bitset_t* bitset, int pg, int offset,
                                       int bit)
{
    while (pg < bitset->npages) {
        while (offset < BITSET_PAGE_ITEMS) {
            bitset_page_data_t val = (bitset->pages[pg]->data[offset] >> bit)
                                   & (((~((bitset_page_data_t)0)) >> bit));
            if (val != 0) {
                int bit1 = FFSL(val);

                if (bit1 != 0) {
                    return BIT_IDX(pg, offset, bit + bit1 - 1);
                }
            }

            offset++;
            bit = 0;
        }

        pg++;
        offset = 0;
        bit    = 0;
    }

    return BIT_IDX(bitset->npages, 0, 0);
}

static int bitset_iter_next(iterator_t* iter, void** data)
{
    unsigned long    pos    = 0;
    iterator_data_t* p      = (iterator_data_t*)iter->data;
    bitset_t*        bitset = RTTI_OBJ_TO_DATA(bitset_t, iter->obj);

    if ((p == NULL) || (bitset == NULL)) {
        return -EINVAL;
    }

    if (p->pg >= bitset->npages) {
        return RTTI_ITER_NO_DATA;
    }

    switch (p->overbit) {
    case BITSET_ITER_OVER_0:
        pos = bitset_next_over0(bitset, p->pg, p->offset, p->bit);
        break;
    case BITSET_ITER_OVER_1:
        pos = bitset_next_over1(bitset, p->pg, p->offset, p->bit);
        break;
    default:
        return -EINVAL;
    }

    if (BIT_PAGE(pos) >= bitset->npages) {
        /* exceed the last page */
        p->pg     = bitset->npages;
        p->offset = 0;
        p->bit    = 0;

        return RTTI_ITER_NO_DATA;
    }

    p->pg     = BIT_PAGE(pos + 1);
    p->offset = BIT_OFFSET(pos + 1);
    p->bit    = BIT_BIT(pos + 1);

    *data = (void*)pos;

    return RTTI_ITER_OK;
}

static int bitset_iter_reset(iterator_t* iter)
{
    iterator_data_t* p = (iterator_data_t*)iter->data;

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

    p->pg     = 0;
    p->offset = 0;
    p->bit    = 0;

    return 1;
}

static void bitset_iter_destroy(iterator_t* iter)
{
    iterator_data_t* p = (iterator_data_t*)iter->data;

    FREEUP(p);

    return;
}

typeinfo_t bitset_type = {
    .name        = "bitset",
    .privateSize = 0,
    .dataSize    = (datasize_fn_t)bitset_datasize,
    .init        = (obj_init_fn_t)bitset_init,
    .destroy     = (obj_destroy_fn_t)bitset_destroy,
    .interfaces
    = { .bitmap = &(intf_bitmap_t) { .test   = (bit_test_fn_t)bitset_test,
                                     .set    = (bit_set_fn_t)bitset_set,
                                     .clr    = (bit_clr_fn_t)bitset_clr,
                                     .clrall = (bit_clrall_fn_t)bitset_clrall },
       .iter
        = &(intf_iterator_t) { .make  = (iter_make_fn_t)bitset_iter_make,
                               .next  = (iter_next_fn_t)bitset_iter_next,
                               .reset = (iter_reset_fn_t)bitset_iter_reset,
                               .destroy
                               = (iter_destroy_fn_t)bitset_iter_destroy } }
};