#include <zinnia/alloc.h>
#include <zinnia/radix.h>

static bool covers(const struct radix_node* node, uint64_t index) {
    return node->shift + RADIX_BITS >= 64 || (index >> (node->shift + RADIX_BITS)) == 0;
}

static size_t slot_of(const struct radix_node* node, uint64_t index) {
    return (index >> node->shift) & (RADIX_SLOTS - 1);
}

static struct radix_node* node_alloc(uint8_t shift, struct radix_node* parent, uint8_t offset) {
    struct radix_node* node = mem_alloc(sizeof(struct radix_node), 0);
    if (!node)
        return nullptr;
    node->shift = shift;
    node->parent = parent;
    node->offset = offset;
    return node;
}

static void free_subtree(struct radix_node* node) {
    if (node->shift > 0) {
        for (size_t i = 0; i < RADIX_SLOTS; i++) {
            if (node->slots[i])
                free_subtree(node->slots[i]);
        }
    }
    mem_free(node);
}

void radix_init(struct radix_tree* tree) {
    tree->root = nullptr;
}

void radix_destroy(struct radix_tree* tree) {
    if (tree->root)
        free_subtree(tree->root);
    tree->root = nullptr;
}

bool radix_empty(struct radix_tree* tree) {
    return tree->root == nullptr;
}

static struct radix_node* leaf_of(struct radix_tree* tree, uint64_t index) {
    struct radix_node* node = tree->root;
    if (!node || !covers(node, index))
        return nullptr;
    while (node->shift > 0) {
        node = node->slots[slot_of(node, index)];
        if (!node)
            return nullptr;
    }
    return node;
}

void* radix_lookup(struct radix_tree* tree, uint64_t index) {
    struct radix_node* leaf = leaf_of(tree, index);
    return leaf ? leaf->slots[slot_of(leaf, index)] : nullptr;
}

void** radix_slot(struct radix_tree* tree, uint64_t index) {
    struct radix_node* leaf = leaf_of(tree, index);
    return leaf ? &leaf->slots[slot_of(leaf, index)] : nullptr;
}

static errno_t grow(struct radix_tree* tree, uint64_t index) {
    if (!tree->root) {
        uint8_t shift = 0;
        while (shift + RADIX_BITS < 64 && (index >> (shift + RADIX_BITS)) != 0)
            shift += RADIX_BITS;
        tree->root = node_alloc(shift, nullptr, 0);
        return tree->root ? 0 : ENOMEM;
    }

    while (!covers(tree->root, index)) {
        struct radix_node* old = tree->root;
        struct radix_node* top = node_alloc(old->shift + RADIX_BITS, nullptr, 0);
        if (!top)
            return ENOMEM;
        top->slots[0] = old;
        top->count = 1;
        for (size_t t = 0; t < RADIX_NUM_TAGS; t++) {
            if (old->tags[t])
                top->tags[t] = 1;
        }
        old->parent = top;
        old->offset = 0;
        tree->root = top;
    }
    return 0;
}

errno_t radix_insert(struct radix_tree* tree, uint64_t index, void* value) {
    errno_t status = grow(tree, index);
    if (status)
        return status;

    struct radix_node* node = tree->root;
    while (node->shift > 0) {
        const size_t slot = slot_of(node, index);
        struct radix_node* child = node->slots[slot];
        if (!child) {
            child = node_alloc(node->shift - RADIX_BITS, node, slot);
            if (!child)
                return ENOMEM;
            node->slots[slot] = child;
            node->count++;
        }
        node = child;
    }

    const size_t slot = slot_of(node, index);
    if (node->slots[slot])
        return EEXIST;
    node->slots[slot] = value;
    node->count++;
    return 0;
}

static void clear_tag_upwards(struct radix_node* node, size_t slot, enum radix_tag tag) {
    while (node) {
        node->tags[tag] &= ~((uint64_t)1 << slot);
        if (node->tags[tag])
            return;
        slot = node->offset;
        node = node->parent;
    }
}

void* radix_remove(struct radix_tree* tree, uint64_t index) {
    struct radix_node* node = leaf_of(tree, index);
    if (!node)
        return nullptr;

    size_t slot = slot_of(node, index);
    void* value = node->slots[slot];
    if (!value)
        return nullptr;

    for (size_t t = 0; t < RADIX_NUM_TAGS; t++) {
        if (node->tags[t] & ((uint64_t)1 << slot))
            clear_tag_upwards(node, slot, t);
    }

    node->slots[slot] = nullptr;
    node->count--;
    while (node->count == 0) {
        struct radix_node* parent = node->parent;
        if (!parent) {
            tree->root = nullptr;
            mem_free(node);
            break;
        }
        slot = node->offset;
        parent->slots[slot] = nullptr;
        parent->count--;
        mem_free(node);
        node = parent;
    }
    return value;
}

void radix_tag_set(struct radix_tree* tree, uint64_t index, enum radix_tag tag) {
    struct radix_node* node = leaf_of(tree, index);
    if (!node || !node->slots[slot_of(node, index)])
        return;
    size_t slot = slot_of(node, index);
    while (node) {
        const uint64_t bit = (uint64_t)1 << slot;
        if (node->tags[tag] & bit)
            return;
        node->tags[tag] |= bit;
        slot = node->offset;
        node = node->parent;
    }
}

void radix_tag_clear(struct radix_tree* tree, uint64_t index, enum radix_tag tag) {
    struct radix_node* node = leaf_of(tree, index);
    if (!node)
        return;
    const size_t slot = slot_of(node, index);
    if (node->tags[tag] & ((uint64_t)1 << slot))
        clear_tag_upwards(node, slot, tag);
}

bool radix_tag_get(struct radix_tree* tree, uint64_t index, enum radix_tag tag) {
    struct radix_node* node = leaf_of(tree, index);
    return node && (node->tags[tag] & ((uint64_t)1 << slot_of(node, index)));
}

bool radix_tagged(struct radix_tree* tree, enum radix_tag tag) {
    return tree->root && tree->root->tags[tag];
}

struct gang {
    size_t max;
    size_t found;
    void** out;
    uint64_t* out_index;
    int tag;
};

static void gang_walk(struct radix_node* node, uint64_t base, uint64_t first, struct gang* g) {
    size_t start = 0;
    if (first > base) {
        const uint64_t delta = (first - base) >> node->shift;
        if (delta >= RADIX_SLOTS)
            return;
        start = delta;
    }

    for (size_t i = start; i < RADIX_SLOTS && g->found < g->max; i++) {
        if (!node->slots[i])
            continue;
        if (g->tag >= 0 && !(node->tags[g->tag] & ((uint64_t)1 << i)))
            continue;
        const uint64_t index = base + ((uint64_t)i << node->shift);
        if (node->shift == 0) {
            if (g->out)
                g->out[g->found] = node->slots[i];
            if (g->out_index)
                g->out_index[g->found] = index;
            g->found++;
        } else {
            gang_walk(node->slots[i], index, first > index ? first : index, g);
        }
    }
}

static size_t gang_lookup(struct radix_tree* tree, uint64_t first, int tag, size_t max, void** out, uint64_t* idx) {
    struct radix_node* root = tree->root;
    if (!root || max == 0 || !covers(root, first))
        return 0;
    struct gang g = {.max = max, .out = out, .out_index = idx, .tag = tag};
    gang_walk(root, 0, first, &g);
    return g.found;
}

size_t radix_gang_lookup(struct radix_tree* tree, uint64_t first, size_t max, void** out, uint64_t* out_index) {
    return gang_lookup(tree, first, -1, max, out, out_index);
}

size_t radix_gang_lookup_tag(
    struct radix_tree* tree,
    uint64_t first,
    enum radix_tag tag,
    size_t max,
    void** out,
    uint64_t* out_index
) {
    return gang_lookup(tree, first, tag, max, out, out_index);
}