#include "c_LazyPrimMST.h"
#include "c_MinPQ.h"


/* ================================================================================================================== */
/* Private Comparison Callback for c_MinPQ_t */

static int c_LazyPrim_EdgeCompare(const void* a, const void* b, void* args) {
    c_size_t id_a = *(const c_size_t*)a;
    c_size_t id_b = *(const c_size_t*)b;
    c_Edge_t* pool = (c_Edge_t*)args;

    double weight_a = pool[id_a].weight;
    double weight_b = pool[id_b].weight;

    return (weight_a > weight_b) - (weight_a < weight_b);
}

/* ================================================================================================================== */
/* Lazy Prim Subroutine Processing Functions */

static void c_LazyPrim_Scan(c_LazyPrimMST_t* self, const c_EdgeWeightedGraph_t* graph, c_size_t v, c_MinPQ_t* pq) {
    self->marked[v] = C_TRUE;

    c_UIntArray_t* adj = &graph->adj_list[v];
    c_size_t size = (c_size_t)c_UIntArray_GetSize(adj);

    for (c_size_t i = 0; i < size; ++i) {
        c_uint_t generic_edge_id = 0;
        c_err_t err = c_UIntArray_Get(adj, i, &generic_edge_id);

        if (err == C_SUCCESS) {
            c_size_t edge_id = (c_size_t)generic_edge_id;
            c_Edge_t* edge = &graph->edges_pool[edge_id];
            c_size_t w = (edge->v == v) ? edge->w : edge->v;

            if (!self->marked[w]) {
                /* Push the edge_id payload into the min-priority queue */
                c_MinPQ_Push(pq, &edge_id);
            }
        }
    }
}

c_err_t c_LazyPrimMST_Init(c_LazyPrimMST_t* self, c_EdgeWeightedGraph_t* graph, c_Allocator_t* allocator) {
    if (!self || !graph) return C_ERR_PARAM;

    self->allocator = allocator ? *allocator : c_DefaultAllocator;
    self->weight = 0.0;
    c_VertexIdList_Init(&self->mst_edges, 0, allocator);

    self->marked = (c_bool_t*)c_Allocator_Calloc(&self->allocator, graph->V, sizeof(c_bool_t));
    if (!self->marked) return C_ERR_NOMEM;

    /* 1. Initialize your formal generic Min-Priority Queue to hold c_size_t elements */
    c_MinPQ_t pq;
    c_err_t err = c_MinPQ_Init(
        &pq,
        graph->E,
        sizeof(c_size_t),
        c_LazyPrim_EdgeCompare,
        graph->edges_pool, /* Pass global edge pool as the comparison context */
        allocator
    );

    if (err != C_SUCCESS) {
        c_Allocator_Free(&self->allocator, self->marked);
        self->marked = NULL;
        return err;
    }

    /* 2. Run over all vertices to build component spanning forest loops if graph is disconnected */
    for (c_size_t v = 0; v < graph->V; ++v) {
        if (!self->marked[v]) {
            c_LazyPrim_Scan(self, graph, v, &pq);

            while (!c_MinPQ_IsEmpty(&pq)) {
                c_size_t edge_id = 0;
                c_MinPQ_Pop(&pq, &edge_id); /* Pops the edge with minimum weight */

                c_Edge_t* edge = &graph->edges_pool[edge_id];
                c_size_t v_end = edge->v;
                c_size_t w_end = edge->w;

                /* Invariant check: Skip if both vertices are already part of the MST tree */
                if (self->marked[v_end] && self->marked[w_end]) continue;

                c_VertexIdList_Append(&self->mst_edges, (c_uint_t)edge_id);
                self->weight += edge->weight;

                if (!self->marked[v_end]) c_LazyPrim_Scan(self, graph, v_end, &pq);
                if (!self->marked[w_end]) c_LazyPrim_Scan(self, graph, w_end, &pq);
            }
        }
    }

    /* Clean up PQ tracking allocations */
    c_MinPQ_Destroy(&pq);
    return C_SUCCESS;
}

void c_LazyPrimMST_Destroy(c_LazyPrimMST_t* self) {
    if (!self) return;
    if (self->marked) {
        c_Allocator_Free(&self->allocator, self->marked);
        self->marked = NULL;
    }
    c_VertexIdList_Destroy(&self->mst_edges);
    self->weight = 0.0;
}

c_err_t c_LazyPrimMST_GetEdges(c_LazyPrimMST_t* self, c_VertexIdList_t* out_edges) {
    if (!self || !out_edges) return C_ERR_PARAM;
    return c_VertexIdList_Copy(out_edges, &self->mst_edges);
}