package dijkstra

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

// Edge is a directed, weighted edge to another vertex.
type Edge struct {
	To     int
	Weight int
}

// Dijkstra returns the shortest distance from source to every vertex.
// Unreachable vertices have distance math.MaxInt.
func Dijkstra(graph [][]Edge, source int) ([]int, error) {
	if source < 0 || source >= len(graph) {
		return nil, fmt.Errorf("source vertex %d is out of range", source)
	}

	for from, edges := range graph {
		for _, edge := range edges {
			if edge.To < 0 || edge.To >= len(graph) {
				return nil, fmt.Errorf("edge %d -> %d is out of range", from, edge.To)
			}
			if edge.Weight < 0 {
				return nil, fmt.Errorf("edge %d -> %d has negative weight", from, edge.To)
			}
		}
	}

	dist := make([]int, len(graph))
	for i := range dist {
		dist[i] = math.MaxInt
	}
	dist[source] = 0

	queue := priorityQueue{{vertex: source, distance: 0}}
	heap.Init(&queue)

	for queue.Len() > 0 {
		current := heap.Pop(&queue).(item)
		if current.distance != dist[current.vertex] {
			continue // Ignore stale queue entries.
		}

		for _, edge := range graph[current.vertex] {
			if current.distance > math.MaxInt-edge.Weight {
				continue // The path length cannot be represented as an int.
			}

			candidate := current.distance + edge.Weight
			if candidate < dist[edge.To] {
				dist[edge.To] = candidate
				heap.Push(&queue, item{vertex: edge.To, distance: candidate})
			}
		}
	}

	return dist, nil
}

type item struct {
	vertex   int
	distance int
}

type priorityQueue []item

func (q priorityQueue) Len() int           { return len(q) }
func (q priorityQueue) Less(i, j int) bool { return q[i].distance < q[j].distance }
func (q priorityQueue) Swap(i, j int)      { q[i], q[j] = q[j], q[i] }

func (q *priorityQueue) Push(value any) {
	*q = append(*q, value.(item))
}

func (q *priorityQueue) Pop() any {
	old := *q
	last := old[len(old)-1]
	*q = old[:len(old)-1]
	return last
}

func TestDijkstra(t *testing.T) {
	graph := [][]Edge{
		{{To: 1, Weight: 4}, {To: 2, Weight: 1}},
		{{To: 3, Weight: 1}},
		{{To: 1, Weight: 2}, {To: 3, Weight: 5}},
		nil,
		nil,
	}

	got, err := Dijkstra(graph, 0)
	if err != nil {
		t.Fatal(err)
	}

	want := []int{0, 3, 1, 4, math.MaxInt}
	if !reflect.DeepEqual(got, want) {
		t.Fatalf("Dijkstra() = %v, want %v", got, want)
	}
}

func TestDijkstraRejectsNegativeWeights(t *testing.T) {
	graph := [][]Edge{
		{{To: 1, Weight: -1}},
		nil,
	}

	if _, err := Dijkstra(graph, 0); err == nil {
		t.Fatal("Dijkstra() accepted a negative edge weight")
	}
}