#include <c_IndexMinPQ.h>
#include <c_Memory.h>

/**
 * Internal helper to swap two positions inside the heap structures.
 * Also keeps the inverted index (qp) tightly synchronized.
 */
C_STATIC_FORCE_INLINE
void c_IndexMinPQ_Swap(c_IndexMinPQ_t* pq_inst, c_size_t i, c_size_t j) {
    long long temp_pq = pq_inst->pq[i];
    pq_inst->pq[i] = pq_inst->pq[j];
    pq_inst->pq[j] = temp_pq;

    pq_inst->qp[pq_inst->pq[i]] = (long long)i;
    pq_inst->qp[pq_inst->pq[j]] = (long long)j;
}

/**
 * Sift-Up Operational Core
 */
C_STATIC_FORCE_INLINE
void c_IndexMinPQ_SiftUp(c_IndexMinPQ_t* pq_inst, c_size_t current) {
    char* keys_arr = (char*)pq_inst->keys;
    c_size_t es = pq_inst->element_size;

    while (current > 0) {
        c_size_t parent = (current - 1) / 2;

        const void* current_key = keys_arr + (pq_inst->pq[current] * es);
        const void* parent_key = keys_arr + (pq_inst->pq[parent] * es);

        if (pq_inst->compar(current_key, parent_key) >= 0) {
            break;
        }
        c_IndexMinPQ_Swap(pq_inst, current, parent);
        current = parent;
    }
}

/**
 * Sift-Down Operational Core
 */
C_STATIC_FORCE_INLINE
void c_IndexMinPQ_SiftDown(c_IndexMinPQ_t* pq_inst, c_size_t current) {
    char* keys_arr = (char*)pq_inst->keys;
    c_size_t es = pq_inst->element_size;

    while (1) {
        c_size_t left_child = (2 * current) + 1;
        c_size_t right_child = (2 * current) + 2;
        c_size_t smallest = current;

        if (left_child < pq_inst->size) {
            if (pq_inst->compar(keys_arr + (pq_inst->pq[left_child] * es), keys_arr + (pq_inst->pq[smallest] * es)) < 0) {
                smallest = left_child;
            }
        }
        if (right_child < pq_inst->size) {
            if (pq_inst->compar(keys_arr + (pq_inst->pq[right_child] * es), keys_arr + (pq_inst->pq[smallest] * es)) < 0) {
                smallest = right_child;
            }
        }

        if (smallest == current) {
            break;
        }
        c_IndexMinPQ_Swap(pq_inst, current, smallest);
        current = smallest;
    }
}

/* ------------------------------------------------------------------------------------------------------------------ */
/*  */

c_err_t c_IndexMinPQ_Init(c_IndexMinPQ_t* pq_inst, c_size_t max_items, c_size_t element_size,
                          int (*compar)(const void*, const void*)) {
    if (pq_inst == NULL || max_items == 0 || element_size == 0 || compar == NULL) return C_ERR_PARAM;

    pq_inst->max_items = max_items;
    pq_inst->element_size = element_size;
    pq_inst->size = 0;
    pq_inst->compar = compar;

    pq_inst->keys = C_ALLOC(max_items * element_size);
    pq_inst->pq = (long long*)C_ALLOC(max_items * sizeof(long long));
    pq_inst->qp = (long long*)C_ALLOC(max_items * sizeof(long long));

    if (pq_inst->keys == NULL || pq_inst->pq == NULL || pq_inst->qp == NULL) {
        // Safe cleanup if any allocation segment drops offline
        C_FREE(pq_inst->keys); C_FREE(pq_inst->pq); C_FREE(pq_inst->qp);
        return C_ERR_NOMEM;
    }

    // Initialize inverted tracking array cells to -1 (indicating absent from heap)
    for (c_size_t i = 0; i < max_items; i++) {
        pq_inst->qp[i] = -1;
    }

    return C_ERR_OK;
}

void c_IndexMinPQ_Destroy(c_IndexMinPQ_t* pq_inst) {
    if (pq_inst) {
        C_FREE(pq_inst->keys); pq_inst->keys = NULL;
        C_FREE(pq_inst->pq);   pq_inst->pq = NULL;
        C_FREE(pq_inst->qp);   pq_inst->qp = NULL;
        pq_inst->size = 0;
        pq_inst->max_items = 0;
    }
}

c_bool_t c_IndexMinPQ_Contains(c_IndexMinPQ_t* pq_inst, c_size_t ext_id) {
    if (pq_inst == NULL || ext_id >= pq_inst->max_items) return C_FALSE;
    return pq_inst->qp[ext_id] != -1;
}

c_err_t c_IndexMinPQ_Push(c_IndexMinPQ_t* pq_inst, c_size_t ext_id, const void* element) {
    if (pq_inst == NULL || ext_id >= pq_inst->max_items || element == NULL) return C_ERR_PARAM;
    if (c_IndexMinPQ_Contains(pq_inst, ext_id)) return C_ERR_ALREADY_EXISTS; // Must use Change Key API if already existing

    char* keys_arr = (char*)pq_inst->keys;

    // Store key inside the primary index table slot
    memcpy(keys_arr + (ext_id * pq_inst->element_size), element, pq_inst->element_size);

    // Append to bottom leaf array trackers
    c_size_t current = pq_inst->size;
    pq_inst->pq[current] = (long long)ext_id;
    pq_inst->qp[ext_id] = (long long)current;

    pq_inst->size++;

    // Sift upward to re-stabilize the min-heap property
    c_IndexMinPQ_SiftUp(pq_inst, current);

    return C_ERR_OK;
}

c_err_t c_IndexMinPQ_Pop(c_IndexMinPQ_t* pq_inst, c_size_t* out_ext_id, void* out_element_buffer) {
    if (pq_inst == NULL) return C_ERR_PARAM;
    if (pq_inst->size == 0) return C_ERR_EMPTY;

    long long min_ext_id = pq_inst->pq[0];
    if (out_ext_id) *out_ext_id = (c_size_t)min_ext_id;

    if (out_element_buffer) {
        char* keys_arr = (char*)pq_inst->keys;
        memcpy(out_element_buffer, keys_arr + (min_ext_id * pq_inst->element_size), pq_inst->element_size);
    }

    // Swap top root with the trailing active leaf node
    c_IndexMinPQ_Swap(pq_inst, 0, pq_inst->size - 1);

    // Clean up tracking registers for extracted ID
    pq_inst->qp[min_ext_id] = -1;
    pq_inst->size--;

    // Sift down from the root to re-balance the tree bounds
    if (pq_inst->size > 0) {
        c_IndexMinPQ_SiftDown(pq_inst, 0);
    }

    return C_ERR_OK;
}

c_err_t c_IndexMinPQ_ChangeKey(c_IndexMinPQ_t* pq_inst, c_size_t ext_id, const void* new_element) {
    if (pq_inst == NULL || ext_id >= pq_inst->max_items || new_element == NULL) return C_ERR_PARAM;
    if (!c_IndexMinPQ_Contains(pq_inst, ext_id)) return C_ERR_NOTFOUND;

    char* keys_arr = (char*)pq_inst->keys;
    c_size_t es = pq_inst->element_size;

    // Update raw payload data in-place
    memcpy(keys_arr + (ext_id * es), new_element, es);

    // Leverage the inverted index (qp) to instantly locate the item's position inside the heap tree
    c_size_t heap_pos = (c_size_t)pq_inst->qp[ext_id];

    // Trigger localized sifting in both directions. Only one will execute depending on whether the key grew or shrank.
    c_IndexMinPQ_SiftUp(pq_inst, heap_pos);
    c_IndexMinPQ_SiftDown(pq_inst, heap_pos);

    return C_ERR_OK;
}

void* c_IndexMinPQ_Peek(c_IndexMinPQ_t* pq_inst, c_size_t* out_ext_id) {
    if (pq_inst == NULL || pq_inst->size == 0) return NULL;
    if (out_ext_id) *out_ext_id = (c_size_t)pq_inst->pq[0];

    char* keys_arr = (char*)pq_inst->keys;
    return keys_arr + (pq_inst->pq[0] * pq_inst->element_size);
}