use std::cmp::Reverse;
use std::collections::BinaryHeap;

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub struct Edge {
    pub to: usize,
    pub cost: u64,
}

/// Computes shortest distances from `start` using Dijkstra's algorithm.
///
/// The graph is an adjacency list: `graph[u]` contains all outgoing edges from `u`.
/// Edge costs must be non-negative, which `u64` enforces.
pub fn dijkstra(graph: &[Vec<Edge>], start: usize) -> (Vec<Option<u64>>, Vec<Option<usize>>) {
    let n = graph.len();
    let mut dist = vec![None; n];
    let mut prev = vec![None; n];

    if start >= n {
        return (dist, prev);
    }

    let mut heap = BinaryHeap::new();
    dist[start] = Some(0);
    heap.push(Reverse((0_u64, start)));

    while let Some(Reverse((cost, node))) = heap.pop() {
        if dist[node] != Some(cost) {
            continue;
        }

        for &Edge { to, cost: edge_cost } in &graph[node] {
            assert!(to < n, "edge points to a node outside the graph");

            let Some(next_cost) = cost.checked_add(edge_cost) else {
                continue;
            };

            if dist[to].map_or(true, |known| next_cost < known) {
                dist[to] = Some(next_cost);
                prev[to] = Some(node);
                heap.push(Reverse((next_cost, to)));
            }
        }
    }

    (dist, prev)
}

/// Reconstructs a shortest path from `start` to `goal` using Dijkstra's predecessor list.
pub fn shortest_path(prev: &[Option<usize>], start: usize, goal: usize) -> Option<Vec<usize>> {
    if start >= prev.len() || goal >= prev.len() {
        return None;
    }

    let mut path = Vec::new();
    let mut current = goal;

    loop {
        path.push(current);

        if current == start {
            path.reverse();
            return Some(path);
        }

        current = prev[current]?;
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn finds_shortest_distances() {
        let graph = vec![
            vec![Edge { to: 1, cost: 4 }, Edge { to: 2, cost: 1 }],
            vec![Edge { to: 3, cost: 1 }],
            vec![Edge { to: 1, cost: 2 }, Edge { to: 3, cost: 5 }],
            vec![],
        ];

        let (dist, _) = dijkstra(&graph, 0);

        assert_eq!(dist, vec![Some(0), Some(3), Some(1), Some(4)]);
    }

    #[test]
    fn reconstructs_shortest_path() {
        let graph = vec![
            vec![Edge { to: 1, cost: 4 }, Edge { to: 2, cost: 1 }],
            vec![Edge { to: 3, cost: 1 }],
            vec![Edge { to: 1, cost: 2 }, Edge { to: 3, cost: 5 }],
            vec![],
        ];

        let (_, prev) = dijkstra(&graph, 0);

        assert_eq!(shortest_path(&prev, 0, 3), Some(vec![0, 2, 1, 3]));
    }

    #[test]
    fn leaves_unreachable_nodes_without_distance() {
        let graph = vec![
            vec![Edge { to: 1, cost: 7 }],
            vec![],
            vec![],
        ];

        let (dist, prev) = dijkstra(&graph, 0);

        assert_eq!(dist, vec![Some(0), Some(7), None]);
        assert_eq!(shortest_path(&prev, 0, 2), None);
    }
}