// Updated timeout_heap.c

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

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

// Forward declaration of timeout_node (from u_async.h, but included here for type safety)
struct timeout_node;

// Private helpers
static void update_index(void *data, size_t idx) {
    if (data) {
        ((struct timeout_node *)data)->heap_index = idx;
    }
}

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;

        // Update indices
        update_index(h->heap[PARENT(i) - 1].data, PARENT(i) - 1);
        update_index(h->heap[i - 1].data, i - 1);

        i = PARENT(i);
    }
}

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;

        // Update indices
        update_index(h->heap[smallest - 1].data, smallest - 1);
        update_index(h->heap[i - 1].data, i - 1);

        i = smallest;
    }
}

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;
    return h;
}

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

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;

    // Set initial index
    update_index(data, idx);

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

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

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

    // Update index for new root
    update_index(h->heap[0].data, 0);

    // 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;
    *out = h->heap[0];
    return 0;
}

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

    DEBUG_DEBUG(DEBUG_CATEGORY_TIMERS, "timeout_heap_pop: entering, size=%zu", h->size);

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

    DEBUG_DEBUG(DEBUG_CATEGORY_TIMERS, "timeout_heap_pop: returning element, data=%p", out->data);
    return 0;
}

int timeout_heap_remove(TimeoutHeap *h, void *data) {
    struct timeout_node *node = (struct timeout_node *)data;
    size_t idx = node->heap_index;
    if (idx == SIZE_MAX || idx >= h->size || h->heap[idx].data != data) {
        return -1;  // Not found or invalid
    }

    // Swap with last
    size_t last = --h->size;
    if (idx != last) {
        TimeoutEntry temp = h->heap[idx];
        h->heap[idx] = h->heap[last];
        h->heap[last] = temp;

        // Update indices
        update_index(h->heap[idx].data, idx);
        update_index(h->heap[last].data, last);

        // Heapify down and up to restore property
        heapify_down(h, idx + 1);
        bubble_up(h, idx + 1);
    }

    // Invalidate index
    node->heap_index = SIZE_MAX;
    return 0;
}