/* text-store - A store for storing reference-counted, de-duplicated strings.
 * Copyright (C) 2025 Jeremia Dominguez <me@jeremia.dev>
 * 
 * This program 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 3 of the License, or
 * (at your option) any later version.
 * 
 * This program 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.
 * 
 * You should have received a copy of the GNU Lesser General Public License
 * along with this program.  If not, see <https://www.gnu.org/licenses/>.
 */

#include <alloca.h>
#include <stdalign.h>
#include <stdbool.h>
#include <stddef.h>
#include <stdint.h>
#include <stdlib.h>
#include <string.h>

#include <textstore.h>

#include "rapidhash.h"

/**
 * TODO:    Make all of this easier to follow.
 *        In hindsight, I realize how much all of this relies on specific pointer math and such.
 *        It may be better to split some things into functions, macros or structs?
 *        This may be helped by splitting the free-queue into its own allocation, rather than
 *        sharing space with the text buffer?
 */

static inline textstore_handle to_handle(uint32_t index, uint32_t generation) {
    return ((uint64_t)generation << 32) | index;
}

/**
 * The hashmap is basically a linked list. The first 32 bits of a map entry are the index part,
 * which is an index of the string table entries. The latter 32 bits are a link to the next map
 * entry, for when a position in the map is occupied.
**/

#define map_index(idx) ((uint32_t)idx)
#define map_link(idx) ((uint32_t)((uint64_t)idx >> 32))
#define map_entry(idx, link) ((uint64_t)link << 32 | idx)
#define map_no_link 0xFFFFFFFF

/**
 * Allocates enough room on the SoA and hash map for n items, and sets the store's pointers.
 *
 * Note: Use on a temporary copy.
 */
static inline void textstore_alloc(textstore_t *store, size_t n) {
    #define X(type, ...) + sizeof(type)
        /* Plus the size of map. */
        const size_t soa_size = sizeof(uint64_t) TEXTSTORE_SOA;
    #undef X

    store->map = (uint64_t *)calloc(n, soa_size);
    store->key = (textstore_slice *)(store->map + n);
    store->ref = (uint32_t *)(store->key + n);
    store->generation = (uint32_t *)(store->ref + n);
    store->active = (bool *)(store->generation + n);
}

static void textstore_put(textstore_t *store,  char *str, uint32_t len, uint32_t value) {
    uint64_t *map     = store->map;
    uint32_t capacity = store->hashmap_cap;
    uint32_t idx      = rapidhash(str, len) % capacity;
    uint32_t last_idx = map_no_link;
start:
    if (map[idx] == 0)
        goto finish;
    if (map_link(map[idx]) != map_no_link) {
        last_idx = idx;
        idx = map_link(map[idx]);
        goto start;
    }
    last_idx = idx;
    while (map[idx] != 0) {
        idx = (idx + 1) % capacity;
    }

finish:
    map[idx] = map_entry(value, map_no_link);
    if (last_idx != map_no_link)
        map[last_idx] = map_entry(map_index(map[last_idx]), idx);
}

static void textstore_rehash(textstore_t *store) {
    char            *buffer = store->buffer;
    textstore_slice *slices = store->key + 1; /* Skip index 0, empty string. */
    bool            *active = store->active + 1;
    uint32_t         count  = store->count - 1;
    memset(store->map, 0, sizeof(uint32_t) * store->hashmap_cap);
    for (uint32_t i = 0; i < count; i++) {
        if (active[i]) {
            textstore_slice slice = slices[i];
            textstore_put(store, &buffer[slice.offset], slice.len, i + 1);
        }
    }
}

static textstore_handle textstore_find(textstore_t *store, char *str, uint32_t len) {
    assert(str != NULL);
    assert(len != 0);

    char            *buffer      = store->buffer;
    uint32_t         hashmap_cap = store->hashmap_cap;
    uint64_t        *map         = store->map;
    textstore_slice *key         = store->key;
    bool            *active      = store->active;

    uint32_t index = rapidhash(str, len) % hashmap_cap;
start:
    if (map[index] == 0)
        goto err;
    uint32_t idx = map_index(map[index]);
    if (key[idx].len == len && active[idx]) {
        if (memcmp(str, &buffer[key[idx].offset], len) == 0)
            return to_handle(idx, store->generation[idx]);
    }
    if (map_link(map[index]) == map_no_link)
        goto err;
    index = map_link(map[index]);
    goto start;
err:
    return TEXTSTORE_HANDLE_ERROR;
}

/**
 * Appends text to the store's buffer.
 */
static void textstore_append(textstore_t *store, char *str, uint32_t len) {
    if (len + 1 + store->length > store->buffer_cap) {
        uint32_t new_capacity= store->buffer_cap * 2;
        while (new_capacity < len + 1 + store->length) {
            new_capacity *= 2;
        }
        store->buffer = realloc(store->buffer, new_capacity);
        memset(&store->buffer[store->buffer_cap], 0, new_capacity - store->buffer_cap);
        store->buffer_cap = new_capacity;
    }
    memcpy(&store->buffer[store->length], str, len);
    store->buffer[store->length + len] = 0;
    store->length += len + 1;
}

textstore_handle textstore_intern(textstore_t *store, char *str) {
    return textstore_intern_len(store, str, strlen(str));
}

textstore_handle textstore_intern_len(textstore_t *store, char *str, uint32_t len) {
    textstore_handle handle = to_handle(0, 0);
    uint32_t index;

    if (len == 0) { /* Special case empty string. */
        textstore_retain(store, handle);
        return handle;
    }
    if (str == NULL) return TEXTSTORE_HANDLE_ERROR;

    handle = textstore_find(store, str, len);
    if (handle != TEXTSTORE_HANDLE_ERROR) {
        textstore_retain(store, handle);
        return handle;
    }

    uint32_t offset = store->length;
    textstore_append(store, str, len);


    if (store->available > 0) {
        /* Pop a freed handle index off the end of the queue. */
        index = ((uint32_t *)store->buffer)[--store->available];
        store->generation[index]++;
    } else {
        index = store->count++;
        store->generation[index] = 0;
    }

    textstore_put(store, str, len, index);
    handle = to_handle(index, store->generation[index]);
    store->key[index] = (textstore_slice){offset, len};
    store->ref[index] = 1;
    store->active[index] = true;

    if (index > ((store->hashmap_cap/4)*3)) {
        /* At 3/4 capacity, grow the hash-table. */
        textstore_t copy = *store;
        textstore_alloc(&copy, copy.hashmap_cap * 2);
        #define X(type, member) memcpy(copy.member, store->member, copy.count * sizeof(type));
            TEXTSTORE_SOA
        #undef X
        copy.hashmap_cap *= 2;
        free(store->map);
        *store = copy;
        textstore_rehash(store);
    }

    return handle;
}

int textstore_gc(textstore_t *store) {
    textstore_slice *strings = store->key;
    uint32_t *refs = store->ref;
    bool *actives = store->active;
    uint32_t count = store->count;
    uint32_t available = store->available;

    /**
     * As a note for whoever is unfortunate enough to interact with my code (likely myself).
     * To minimize the number of allocations, the textstore keeps the queue of freed handles
     * and the actual text, in contiguous memory. The free queue is at the start of the buffer,
     * and new strings get added to the end of the buffer. This works out, because the queue list
     * only grows when textstore_gc() is caled, when all handles with 0 references are collected.
     *
     * The free-queue is last-in-first-out.
     *
     * [ 01 ] [ 02 ]   H e l l o   W o r l d ! ...
     *           ^
     *           store->available = 2
     *
     * So when a new string is interned, and there are handles available to use, the last item is
     * popped off the stack.
     * 
     * [ 01 ] [ 02 ]   H e l l o   W o r l d ! ...
     *    ^
     *    store->available = 1
     */

    // Pass 1: Count how many strings need to be collected, and determine how
    //         much space is needed for the new buffer.
    uint32_t collect = 0;
    uint32_t new_length = available * sizeof(uint32_t);
    for (uint32_t i = 0; i < count; i++) {
        if (actives[i]) {
            if (refs[i] > 0) {
                new_length += strings[i].len + 1;
            } else {
                collect += 1;
                new_length += sizeof(uint32_t);
            }
        }
    }

    // If no strings are in need of collecting, stop early.
    if (collect == 0) return 0;
    
    // Pass 2: Generate new buffer from the active strings.
    if (new_length < store->buffer_cap)
        new_length = store->buffer_cap;
    char *new_buffer = malloc(new_length);
    if (new_buffer == NULL) return 1;
    memset(new_buffer, 0, new_length);
    uint32_t new_offset = (collect + available) * sizeof(uint32_t);
    // copy existing free queue to the new one.
    if (available)
        memcpy(new_buffer, store->buffer, available * sizeof(uint32_t));
    uint32_t *queue_buffer = (uint32_t *)new_buffer;
    uint32_t queue_offset = available;

    for (uint32_t i = 0; i < count; i++) {
        if (actives[i]) {
            if (refs[i] == 0) {
                queue_buffer[queue_offset++] = i;
                actives[i] = false;
                strings[i] = (textstore_slice){0};
                continue;
            }
            // Copy string to new location, and shift offsets.
            memcpy(&new_buffer[new_offset], &store->buffer[strings[i].offset], strings[i].len);
            strings[i].offset = new_offset;
            new_offset += strings[i].len + 1;
        }
    }

    free(store->buffer);
    store->buffer = new_buffer;
    store->buffer_cap = new_length;
    store->available = collect + available;

    textstore_rehash(store);
    return 0;
}

int textstore_init(textstore_t *store, uint32_t initial_buffer) {
    textstore_t temp = {0};
    if (initial_buffer == 0) initial_buffer = 1;
    temp.buffer = (char *)calloc(initial_buffer, 1);
    if (temp.buffer == NULL) return 1;
    temp.buffer_cap = initial_buffer;

    temp.hashmap_cap = 8;
    textstore_alloc(&temp, temp.hashmap_cap);

    /* Special-case empty string ("") for index 0. */
    temp.key[0] = (textstore_slice){ 0, 0 };
    temp.ref[0] = 1;
    temp.active[0] = true;
    temp.length++;
    temp.count++;

    *store = temp;
    return 0;
}

void textstore_free(textstore_t *store) {
    free(store->buffer);
    free(store->map);

    *store = (textstore_t){0};
}