package dijkstra

import (
	"container/heap"
	"math"
	"testing"
)

type Edge struct {
	To     string
	Weight int
}

type Graph map[string][]Edge

type node struct {
	name string
	dist int
}

type priorityQueue []node

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

// ShortestPath returns the minimum distance and path from source to target.
func ShortestPath(g Graph, source, target string) (int, []string, bool) {
	dist := map[string]int{source: 0}
	prev := map[string]string{}

	pq := &priorityQueue{{name: source, dist: 0}}
	heap.Init(pq)

	for pq.Len() > 0 {
		cur := heap.Pop(pq).(node)

		// Ignore stale queue entries created before a better path was found.
		if cur.dist != dist[cur.name] {
			continue
		}
		if cur.name == target {
			return cur.dist, buildPath(prev, source, target), true
		}

		for _, e := range g[cur.name] {
			if e.Weight < 0 {
				continue // Dijkstra requires non-negative weights.
			}

			nextDist := cur.dist + e.Weight
			if old, ok := dist[e.To]; !ok || nextDist < old {
				dist[e.To] = nextDist
				prev[e.To] = cur.name
				heap.Push(pq, node{name: e.To, dist: nextDist})
			}
		}
	}

	return math.MaxInt, nil, false
}

func buildPath(prev map[string]string, source, target string) []string {
	path := []string{}
	for at := target; ; at = prev[at] {
		path = append(path, at)
		if at == source {
			break
		}
	}

	for i, j := 0, len(path)-1; i < j; i, j = i+1, j-1 {
		path[i], path[j] = path[j], path[i]
	}
	return path
}

func TestShortestPath(t *testing.T) {
	g := Graph{
		"A": {{To: "B", Weight: 4}, {To: "C", Weight: 2}},
		"B": {{To: "D", Weight: 5}},
		"C": {{To: "B", Weight: 1}, {To: "D", Weight: 8}},
		"D": nil,
	}

	dist, path, ok := ShortestPath(g, "A", "D")
	if !ok {
		t.Fatal("expected a path")
	}
	if dist != 8 {
		t.Fatalf("distance = %d, want 8", dist)
	}

	want := []string{"A", "C", "B", "D"}
	if len(path) != len(want) {
		t.Fatalf("path = %v, want %v", path, want)
	}
	for i := range want {
		if path[i] != want[i] {
			t.Fatalf("path = %v, want %v", path, want)
		}
	}
}

func TestShortestPathUnreachable(t *testing.T) {
	g := Graph{
		"A": {{To: "B", Weight: 1}},
		"C": nil,
	}

	_, path, ok := ShortestPath(g, "A", "C")
	if ok {
		t.Fatal("expected no path")
	}
	if path != nil {
		t.Fatalf("path = %v, want nil", path)
	}
}