package core

import (
	"container/heap"
)

type pqItem struct {
	node   int
	weight float64
	index  int
}

type priorityQueue []*pqItem

func (pq priorityQueue) Len() int { return len(pq) }
func (pq priorityQueue) Less(i, j int) bool {
	return pq[i].weight < pq[j].weight
}
func (pq priorityQueue) Swap(i, j int) {
	pq[i], pq[j] = pq[j], pq[i]
	pq[i].index = i
	pq[j].index = j
}
func (pq *priorityQueue) Push(x any) {
	n := len(*pq)
	item := x.(*pqItem)
	item.index = n
	*pq = append(*pq, item)
}
func (pq *priorityQueue) Pop() any {
	old := *pq
	n := len(old)
	item := old[n-1]
	old[n-1] = nil
	item.index = -1
	*pq = old[0 : n-1]
	return item
}

type PathResult struct {
	TotalWeight float64
	TotalTime   float64
	TotalRisk   float64
	Path        []int
}

// SingleSourceDijkstra computes shortest paths from source to all reachable nodes
func (g *Graph) SingleSourceDijkstra(source int) map[int]*PathResult {
	dist := make(map[int]float64)
	timeDist := make(map[int]float64)
	riskDist := make(map[int]float64)
	prev := make(map[int]int)

	pq := make(priorityQueue, 0)
	heap.Init(&pq)

	dist[source] = 0.0
	heap.Push(&pq, &pqItem{node: source, weight: 0.0})

	for pq.Len() > 0 {
		item := heap.Pop(&pq).(*pqItem)
		u := item.node
		w := item.weight

		// Outdated element
		if d, ok := dist[u]; ok && w > d {
			continue
		}

		for _, edge := range g.Adj[u] {
			v := edge.To
			alt := w + edge.Weight
			if d, ok := dist[v]; !ok || alt < d {
				dist[v] = alt
				timeDist[v] = timeDist[u] + edge.Time
				riskDist[v] = riskDist[u] + edge.Risk
				prev[v] = u
				heap.Push(&pq, &pqItem{node: v, weight: alt})
			}
		}
	}

	results := make(map[int]*PathResult)
	for node, w := range dist {
		// Reconstruct path
		var path []int
		curr := node
		for {
			path = append(path, curr)
			p, hasPrev := prev[curr]
			if !hasPrev {
				break
			}
			curr = p
		}
		
		// Reverse path
		for i, j := 0, len(path)-1; i < j; i, j = i+1, j-1 {
			path[i], path[j] = path[j], path[i]
		}

		results[node] = &PathResult{
			TotalWeight: w,
			TotalTime:   timeDist[node],
			TotalRisk:   riskDist[node],
			Path:        path,
		}
	}

	return results
}