#include <c_DijkstraUndirectedSP.h>
#include "c_IndexMinPQ.h"

#define SP_SENTINEL ((c_size_t)-1)

static int c_DijkstraUndirected_DistCompare(const void* a, const void* b, void* args) {
    double dist_a = *(const double*)a;
    double dist_b = *(const double*)b;
    (void)args;
    return (dist_a > dist_b) - (dist_a < dist_b);
}

static void c_DijkstraUndirected_Relax(c_DijkstraUndirectedSP_t* self, c_EdgeWeightedGraph_t* graph, c_size_t v, c_IndexMinPQ_t* pq) {
    c_AdjList_t* adj = &graph->adj_list[v];
    c_size_t size = (c_size_t)c_AdjList_GetSize(adj);

    for (c_size_t i = 0; i < size; ++i) {
        c_uint_t generic_edge_id = 0;
        c_err_t err = c_AdjList_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];

            /* Resolve the other endpoint of the undirected edge relative to v */
            c_size_t w = (edge->v == v) ? edge->w : edge->v;

            /* Guard against negative edge weights which break Dijkstra's structural invariants */
            if (edge->weight < 0.0) continue;

            if (self->dist_to[w] > self->dist_to[v] + edge->weight) {
                self->dist_to[w] = self->dist_to[v] + edge->weight;
                self->edge_to[w] = edge_id;
                self->from_vertex[w] = v; /* Record vertex node transition step */

                if (c_IndexMinPQ_Contains(pq, w)) {
                    c_IndexMinPQ_Change(pq, w, &self->dist_to[w]);
                } else {
                    c_IndexMinPQ_Push(pq, w, &self->dist_to[w]);
                }
            }
        }
    }
}

c_err_t c_DijkstraUndirectedSP_Init(c_DijkstraUndirectedSP_t* self, c_EdgeWeightedGraph_t* graph, c_size_t s, c_Allocator_t* allocator) {
    if (!self || !graph || s >= graph->V) return C_ERR_PARAM;

    self->allocator = allocator ? *allocator : c_DefaultAllocator;
    self->s = s;
    self->V = graph->V;

    self->edge_to     = (c_size_t*)c_Allocator_Calloc(&self->allocator, self->V, sizeof(c_size_t));
    self->from_vertex = (c_size_t*)c_Allocator_Calloc(&self->allocator, self->V, sizeof(c_size_t));
    self->dist_to     = (double*)c_Allocator_Calloc(&self->allocator, self->V, sizeof(double));

    if (!self->edge_to || !self->from_vertex || !self->dist_to) {
        c_DijkstraUndirectedSP_Destroy(self);
        return C_ERR_NOMEM;
    }

    for (c_size_t v = 0; v < self->V; ++v) {
        self->dist_to[v]     = DBL_MAX;
        self->edge_to[v]     = SP_SENTINEL;
        self->from_vertex[v] = SP_SENTINEL;
    }
    self->dist_to[s] = 0.0;

    c_IndexMinPQ_t pq;
    c_err_t err = c_IndexMinPQ_Init(&pq, self->V, sizeof(double), c_DijkstraUndirected_DistCompare, NULL, allocator);
    if (err != C_SUCCESS) {
        c_DijkstraUndirectedSP_Destroy(self);
        return err;
    }

    c_IndexMinPQ_Push(&pq, s, &self->dist_to[s]);

    while (pq.size > 0) {
        c_size_t v = 0;
        c_IndexMinPQ_Pop(&pq, &v);
        c_DijkstraUndirected_Relax(self, graph, v, &pq);
    }

    c_IndexMinPQ_Destroy(&pq);
    return C_SUCCESS;
}

void c_DijkstraUndirectedSP_Destroy(c_DijkstraUndirectedSP_t* self) {
    if (!self) return;
    if (self->edge_to)     c_Allocator_Free(&self->allocator, self->edge_to);
    if (self->from_vertex) c_Allocator_Free(&self->allocator, self->from_vertex);
    if (self->dist_to)     c_Allocator_Free(&self->allocator, self->dist_to);

    self->edge_to     = NULL;
    self->from_vertex = NULL;
    self->dist_to     = NULL;
    self->V = 0;
    self->s = 0;
}

c_err_t c_DijkstraUndirectedSP_PathTo(c_DijkstraUndirectedSP_t* self, c_size_t v, c_VertexIdList_t* out_path) {
    if (!self || !out_path || v >= self->V) return C_ERR_PARAM;
    if (!c_DijkstraUndirectedSP_HasPathTo(self, v)) return C_ERR_FAIL;

    c_size_t* edge_stack = (c_size_t*)c_Allocator_Calloc(&self->allocator, self->V, sizeof(c_size_t));
    if (!edge_stack) return C_ERR_NOMEM;

    c_size_t stack_size = 0;
    c_size_t curr_v = v;

    /* Trace backward via recorded source transitions */
    while (curr_v != self->s) {
        c_size_t edge_id = self->edge_to[curr_v];
        if (edge_id == SP_SENTINEL) break;

        edge_stack[stack_size++] = edge_id;
        curr_v = self->from_vertex[curr_v];
    }

    /* Flip and append forward into output vertex ID list */
    c_err_t err = C_SUCCESS;
    while (stack_size > 0) {
        c_size_t target_edge_id = edge_stack[--stack_size];
        err = c_VertexIdList_Append(out_path, (c_uint_t)target_edge_id);
        if (err != C_SUCCESS) break;
    }

    c_Allocator_Free(&self->allocator, edge_stack);
    return err;
}