#include <rlang.h>
#include "decl/dyn-list-of-decl.h"

#define R_DYN_LOF_GROWTH_FACTOR 2

#define R_DYN_LOF_INIT_SIZE 32

enum shelter_dyn_list_of {
    SHELTER_DYN_LOF_raw,
    SHELTER_DYN_LOF_reserve,
    SHELTER_DYN_LOF_arr_locs,
    SHELTER_DYN_LOF_extra_array,
    SHELTER_DYN_LOF_extra_shelter_array,
    SHELTER_DYN_LOF_moved_arr,
    SHELTER_DYN_LOF_moved_shelter_arr,
    SHELTER_DYN_LOF_arrays,
    SHELTER_DYN_LOF_SIZE
};

struct r_dyn_list_of* r_new_dyn_list_of(
    enum r_type type,
    r_ssize capacity,
    r_ssize width
) {
    switch (type) {
    case R_TYPE_character:
    case R_TYPE_list:
        r_abort("Can't create a dynamic list of barrier vectors.");
    default:
        break;
    }

    r_obj* shelter = KEEP(r_alloc_list(SHELTER_DYN_LOF_SIZE));

    r_obj* lof_raw = r_alloc_raw(sizeof(struct r_dyn_list_of));
    r_list_poke(shelter, SHELTER_DYN_LOF_raw, lof_raw);

    struct r_dyn_array* p_moved_arr =
        r_new_dyn_array(sizeof(struct r_dyn_array*), R_DYN_LOF_INIT_SIZE);
    r_list_poke(shelter, SHELTER_DYN_LOF_moved_arr, p_moved_arr->shelter);

    struct r_dyn_array* p_moved_shelter_arr =
        r_new_dyn_vector(R_TYPE_list, R_DYN_LOF_INIT_SIZE);
    r_list_poke(
        shelter,
        SHELTER_DYN_LOF_moved_shelter_arr,
        p_moved_shelter_arr->shelter
    );

    r_obj* reserve = r_alloc_vector(type, r_ssize_mult(capacity, width));
    r_list_poke(shelter, SHELTER_DYN_LOF_reserve, reserve);
    void* v_reserve = r_vec_begin(reserve);

    r_obj* arr_locs = r_alloc_raw(sizeof(r_ssize) * capacity);
    r_list_poke(shelter, SHELTER_DYN_LOF_arr_locs, arr_locs);
    r_ssize* v_arr_locs = r_raw_begin(arr_locs);
    R_MEM_SET(r_ssize, v_arr_locs, -1, capacity);

    struct r_dyn_array* p_arrays =
        r_new_dyn_array(sizeof(struct r_pair_ptr_ssize), capacity);
    r_list_poke(shelter, SHELTER_DYN_LOF_arrays, p_arrays->shelter);

    struct r_dyn_list_of* p_lof = r_raw_begin(lof_raw);
    *p_lof = (struct r_dyn_list_of) {
        .shelter = shelter,
        .count = 0,
        .capacity = capacity,
        .growth_factor = R_DYN_LOF_GROWTH_FACTOR,

        .v_data = r_dyn_begin(p_arrays),

        // private:
        .width = width,
        .type = type,
        .elt_byte_size = r_vec_elt_sizeof0(type),

        .reserve = reserve,
        .v_reserve = v_reserve,

        .p_moved_arr = p_moved_arr,
        .p_moved_shelter_arr = p_moved_shelter_arr,

        .arr_locs = arr_locs,
        .v_arr_locs = v_arr_locs,

        .p_arrays = p_arrays,
    };

    FREE(1);
    return p_lof;
}

r_obj* r_lof_unwrap(struct r_dyn_list_of* p_lof) {
    r_obj* out = KEEP(r_alloc_list(p_lof->count));

    enum r_type type = p_lof->type;
    r_ssize n = p_lof->count;
    struct r_pair_ptr_ssize* v_arrays = r_dyn_begin(p_lof->p_arrays);

    for (r_ssize i = 0; i < n; ++i) {
        struct r_pair_ptr_ssize array = v_arrays[i];
        r_list_poke(out, i, r_vec_n(type, array.ptr, array.size));
    }

    FREE(1);
    return out;
}

static void r_lof_resize(struct r_dyn_list_of* p_lof, r_ssize capacity) {
    r_ssize count = p_lof->count;

    // Resize reserve
    r_obj* reserve = r_vec_resize0(
        p_lof->type,
        p_lof->reserve,
        r_ssize_mult(capacity, p_lof->width)
    );
    r_list_poke(p_lof->shelter, SHELTER_DYN_LOF_reserve, reserve);

    p_lof->reserve = reserve;
    p_lof->v_reserve = r_vec_begin0(p_lof->type, reserve);
    p_lof->capacity = capacity;

    // Resize array indirections
    r_obj* arr_locs =
        r_raw_resize(p_lof->arr_locs, r_ssize_mult(sizeof(r_ssize), capacity));
    r_list_poke(p_lof->shelter, SHELTER_DYN_LOF_arr_locs, arr_locs);

    r_ssize* v_arr_locs = r_raw_begin(arr_locs);
    r_ssize n_new = capacity - count;
    R_MEM_SET(r_ssize, v_arr_locs + count, -1, n_new);

    p_lof->arr_locs = arr_locs;
    p_lof->v_arr_locs = v_arr_locs;

    // Resize addresses and update them to point to the new memory
    r_dyn_resize(p_lof->p_arrays, capacity);

    struct r_pair_ptr_ssize* v_data = r_dyn_begin(p_lof->p_arrays);
    p_lof->v_data = v_data;

    unsigned char* v_reserve_u = (unsigned char*) p_lof->v_reserve;
    r_ssize bytes = p_lof->width * p_lof->elt_byte_size;

    for (r_ssize i = 0; i < count; ++i) {
        // Preserve addresses of moved arrays
        if (v_arr_locs[i] < 0) {
            r_ssize offset = i * bytes;
            v_data[i].ptr = v_reserve_u + offset;
        }
    }
}

void r_lof_push_back(struct r_dyn_list_of* p_lof) {
    r_ssize count = p_lof->count + 1;

    if (count > p_lof->capacity) {
        r_ssize new_size =
            r_ssize_mult(p_lof->capacity, R_DYN_LOF_GROWTH_FACTOR);
        r_lof_resize(p_lof, new_size);
    }
    p_lof->count = count;

    unsigned char* v_reserve_u = (unsigned char*) p_lof->v_reserve;
    r_ssize offset = (count - 1) * p_lof->width * p_lof->elt_byte_size;

    struct r_pair_ptr_ssize info = {.ptr = v_reserve_u + offset, .size = 0};
    r_dyn_push_back(p_lof->p_arrays, &info);
}

void r_lof_arr_push_back(struct r_dyn_list_of* p_lof, r_ssize i, void* p_elt) {
    if (i >= p_lof->count) {
        r_stop_internal("Location %d does not exist.", i);
    }

    if (reserve_push_back(p_lof, i, p_elt)) {
        return;
    }

    struct r_dyn_array* p_arr = p_lof->p_moved_arr;
    r_ssize arr_i = p_lof->v_arr_locs[i];

    if (arr_i >= p_arr->count) {
        r_stop_internal("Location %d does not exist in the extra array", arr_i);
    }

    struct r_dyn_array* p_inner_arr =
        R_DYN_GET(struct r_dyn_array*, p_arr, arr_i);
    r_dyn_push_back(p_inner_arr, p_elt);

    // Also update pointer in case of resize
    R_DYN_POKE(
        struct r_pair_ptr_ssize,
        p_lof->p_arrays,
        i,
        ((struct r_pair_ptr_ssize) {.ptr = r_dyn_begin(p_inner_arr),
                                    .size = p_inner_arr->count})
    );
}

static bool reserve_push_back(
    struct r_dyn_list_of* p_lof,
    r_ssize i,
    void* p_elt
) {
    if (p_lof->v_arr_locs[i] >= 0) {
        return false;
    }

    struct r_pair_ptr_ssize* p_arr_info = r_dyn_pointer(p_lof->p_arrays, i);
    if (p_arr_info->size >= p_lof->width) {
        // Inner array is getting too big for the reserve. Move it to a
        // dynamic array.
        reserve_move(p_lof, i, p_elt);
        return false;
    }

    r_ssize count = ++p_arr_info->size;

    r_ssize offset = (i * p_lof->width + count - 1) * p_lof->elt_byte_size;
    void* p = ((unsigned char*) p_lof->v_reserve) + offset;

    if (p_elt) {
        r_memcpy(p, p_elt, p_lof->elt_byte_size);
    } else {
        r_memset(p, 0, p_lof->elt_byte_size);
    }

    return true;
}

static void reserve_move(struct r_dyn_list_of* p_lof, r_ssize i, void* p_elt) {
    struct r_dyn_array* p_moved_arr = p_lof->p_moved_arr;
    r_ssize n = p_lof->width;

    struct r_dyn_array* p_new = r_new_dyn_vector(p_lof->type, p_lof->width);
    r_dyn_list_push_back(p_lof->p_moved_shelter_arr, p_new->shelter);
    r_dyn_push_back(p_moved_arr, &p_new);

    void* v_new = r_dyn_begin(p_new);
    void* v_old = R_DYN_GET(struct r_pair_ptr_ssize, p_lof->p_arrays, i).ptr;
    r_memcpy(v_new, v_old, r_ssize_mult(n, p_lof->elt_byte_size));

    p_new->count = n;

    R_DYN_POKE(
        struct r_pair_ptr_ssize,
        p_lof->p_arrays,
        i,
        ((struct r_pair_ptr_ssize) {.ptr = v_new, .size = n})
    );

    p_lof->v_arr_locs[i] = p_moved_arr->count - 1;
}
