// timeout_heap.c

#include "timeout_heap.h"
#include <stdlib.h>
#include <stdio.h>   // For potential error printing, optional

// Helper macros for 1-based indices
#define PARENT(i) ((i) / 2)
#define LEFT_CHILD(i) (2 * (i))
#define RIGHT_CHILD(i) (2 * (i) + 1)

TimeoutHeap *timeout_heap_create(size_t initial_capacity) {
    TimeoutHeap *h = malloc(sizeof(TimeoutHeap));
    if (!h) return NULL;
    h->heap = malloc(sizeof(TimeoutEntry) * initial_capacity);
    if (!h->heap) {
        free(h);
        return NULL;
    }
    h->size = 0;
    h->capacity = initial_capacity;
    h->freed_count = 0;
    h->user_data = NULL;
    h->free_callback = NULL;
    return h;
}

void timeout_heap_destroy(TimeoutHeap *h) {
    if (h) {
        free(h->heap);
        free(h);
    }
}

static void bubble_up(TimeoutHeap *h, size_t i) {
    // i is 1-based
    while (i > 1 && h->heap[PARENT(i) - 1].expiration > h->heap[i - 1].expiration) {
        // Swap with parent
        TimeoutEntry temp = h->heap[PARENT(i) - 1];
        h->heap[PARENT(i) - 1] = h->heap[i - 1];
        h->heap[i - 1] = temp;
        i = PARENT(i);
    }
}

int timeout_heap_push(TimeoutHeap *h, TimeoutTime expiration, void *data) {
    if (h->size == h->capacity) {
        size_t new_cap = h->capacity ? h->capacity * 2 : 1;
        TimeoutEntry *new_heap = realloc(h->heap, sizeof(TimeoutEntry) * new_cap);
        if (!new_heap) return -1;  // Allocation failed
        h->heap = new_heap;
        h->capacity = new_cap;
    }

    // Insert at end (0-based)
    size_t idx = h->size++;
    h->heap[idx].expiration = expiration;
    h->heap[idx].data = data;
    h->heap[idx].deleted = 0;

    // Bubble up (1-based)
    bubble_up(h, idx + 1);
    return 0;
}

static void heapify_down(TimeoutHeap *h, size_t i) {
    // i is 1-based
    while (1) {
        size_t smallest = i;
        size_t left = LEFT_CHILD(i);
        size_t right = RIGHT_CHILD(i);

        if (left <= h->size && h->heap[left - 1].expiration < h->heap[smallest - 1].expiration) {
            smallest = left;
        }
        if (right <= h->size && h->heap[right - 1].expiration < h->heap[smallest - 1].expiration) {
            smallest = right;
        }
        if (smallest == i) break;

        // Swap
        TimeoutEntry temp = h->heap[smallest - 1];
        h->heap[smallest - 1] = h->heap[i - 1];
        h->heap[i - 1] = temp;
        i = smallest;
    }
}

static void remove_root(TimeoutHeap *h) {
    if (h->size == 0) return;

    // Move last to root
    h->heap[0] = h->heap[--h->size];

    // Heapify down (1-based)
    if (h->size > 0) {
        heapify_down(h, 1);
    }
}

int timeout_heap_peek(TimeoutHeap *h, TimeoutEntry *out) {
    if (h->size == 0) return -1;

    // Skip deleted
    size_t i = 0;
    while (i < h->size && h->heap[0].deleted) {
        remove_root(h);
    }
    if (h->size == 0) return -1;

    *out = h->heap[0];
    return 0;
}

int timeout_heap_pop(TimeoutHeap *h, TimeoutEntry *out) {
    if (h->size == 0) return -1;

    // Skip deleted and free their data
    while (h->size > 0 && h->heap[0].deleted) {
        void *data_to_free = h->heap[0].data;
        remove_root(h);
        if (h->free_callback) {
            h->free_callback(h->user_data, data_to_free);
        } else {
            free(data_to_free);
        }
        h->freed_count++;
    }
    if (h->size == 0) return -1;

    *out = h->heap[0];
    remove_root(h);
    return 0;
}

int timeout_heap_cancel(TimeoutHeap *h, TimeoutTime expiration, void *data) {
    for (size_t i = 0; i < h->size; ++i) {
        if (h->heap[i].expiration == expiration && h->heap[i].data == data) {
            h->heap[i].deleted = 1;
            return 0;
        }
    }
    return -1;  // Not found
}

void timeout_heap_set_free_callback(TimeoutHeap *h, void* user_data, void (*callback)(void* user_data, void* data)) {
    if (!h) return;
    h->user_data = user_data;
    h->free_callback = callback;
}