yog

import gleam/dict
import gleam/list
import gleam/order.{type Order}
import gleam/result
import gleamy/priority_queue
import yog/model.{type Graph, type NodeId}

/// Performs a topological sort on a directed graph using Kahn's algorithm.
///
/// Returns a linear ordering of nodes such that for every directed edge (u, v),
/// node u comes before node v in the ordering.
///
/// Returns `Error(Nil)` if the graph contains a cycle.
///
/// **Time Complexity:** O(V + E) where V is vertices and E is edges
///
/// ## Example
///
/// ```gleam
/// topological_sort.topological_sort(graph)
/// // => Ok([1, 2, 3, 4])  // Valid ordering
/// // or Error(Nil)         // Cycle detected
/// ```
pub fn topological_sort(graph: Graph(n, e)) -> Result(List(NodeId), Nil) {
  // 1. Get all unique NodeIds from both out_edges and in_edges
  let all_nodes = model.all_nodes(graph)

  // 2. Calculate initial in-degrees
  let in_degrees =
    all_nodes
    |> list.map(fn(id) {
      let degree =
        dict.get(graph.in_edges, id)
        |> result.map(dict.size)
        |> result.unwrap(0)
      #(id, degree)
    })
    |> dict.from_list()

  // 3. Find starting nodes (in-degree 0)
  let queue =
    dict.to_list(in_degrees)
    |> list.filter(fn(pair) { pair.1 == 0 })
    |> list.map(fn(pair) { pair.0 })

  do_kahn(graph, queue, in_degrees, [], list.length(all_nodes))
}

/// Performs a topological sort that returns the lexicographically smallest sequence.
///
/// Uses a heap-based version of Kahn's algorithm to ensure that when multiple
/// nodes have in-degree 0, the smallest one (according to `compare_nodes`) is chosen first.
///
/// The comparison function operates on **node data**, not node IDs, allowing intuitive
/// comparisons like `string.compare` for alphabetical ordering.
///
/// Returns `Error(Nil)` if the graph contains a cycle.
///
/// **Time Complexity:** O(V log V + E) due to heap operations
///
/// ## Example
///
/// ```gleam
/// // Get alphabetical ordering by node data
/// topological_sort.lexicographical_topological_sort(graph, string.compare)
/// // => Ok([0, 1, 2])  // Node IDs ordered by their string data
///
/// // Custom comparison by priority
/// topological_sort.lexicographical_topological_sort(graph, fn(a, b) {
///   int.compare(a.priority, b.priority)
/// })
/// ```
pub fn lexicographical_topological_sort(
  graph: Graph(n, e),
  compare_nodes: fn(n, n) -> Order,
) -> Result(List(NodeId), Nil) {
  // 1. Get all nodes from the edge maps
  let all_nodes = model.all_nodes(graph)

  // 2. Initial in-degrees
  let in_degrees =
    all_nodes
    |> list.map(fn(id) {
      let degree =
        dict.get(graph.in_edges, id)
        |> result.map(dict.size)
        |> result.unwrap(0)
      #(id, degree)
    })
    |> dict.from_list()

  // 3. Create a comparison function that compares node data instead of IDs
  let compare_by_data = fn(id_a: NodeId, id_b: NodeId) -> Order {
    case dict.get(graph.nodes, id_a), dict.get(graph.nodes, id_b) {
      Ok(data_a), Ok(data_b) -> compare_nodes(data_a, data_b)
      // Fallback to ID comparison if node data not found (shouldn't happen)
      _, _ -> order.Eq
    }
  }

  // 4. Find initial nodes with 0 in-degree and put them in the priority queue
  let initial_queue =
    dict.to_list(in_degrees)
    |> list.filter(fn(pair) { pair.1 == 0 })
    |> list.map(fn(pair) { pair.0 })
    |> list.fold(priority_queue.new(compare_by_data), fn(q, id) {
      priority_queue.push(q, id)
    })

  do_lexical_kahn(graph, initial_queue, in_degrees, [], list.length(all_nodes))
}

fn do_lexical_kahn(graph, q, in_degrees, acc, total_count) {
  case priority_queue.pop(q) {
    Error(Nil) -> {
      case list.length(acc) == total_count {
        True -> Ok(list.reverse(acc))
        False -> Error(Nil)
      }
    }
    Ok(#(head, rest_q)) -> {
      let neighbors = model.successor_ids(graph, head)

      let #(next_q, next_in_degrees) =
        list.fold(neighbors, #(rest_q, in_degrees), fn(state, neighbor) {
          let #(current_q, degrees) = state
          let current_degree = dict.get(degrees, neighbor) |> result.unwrap(0)
          let new_degree = current_degree - 1

          let new_degrees = dict.insert(degrees, neighbor, new_degree)
          let updated_q = case new_degree == 0 {
            True -> priority_queue.push(current_q, neighbor)
            False -> current_q
          }
          #(updated_q, new_degrees)
        })

      do_lexical_kahn(
        graph,
        next_q,
        next_in_degrees,
        [head, ..acc],
        total_count,
      )
    }
  }
}

fn do_kahn(graph, queue, in_degrees, acc, total_node_count) {
  case queue {
    [] -> {
      case list.length(acc) == total_node_count {
        True -> Ok(list.reverse(acc))
        False -> Error(Nil)
        // Cycle detected!
      }
    }
    [head, ..tail] -> {
      let neighbors = model.successor_ids(graph, head)

      let #(next_queue, next_in_degrees) =
        list.fold(neighbors, #(tail, in_degrees), fn(state, neighbor) {
          let #(q, degrees) = state
          let current_degree = dict.get(degrees, neighbor) |> result.unwrap(0)
          let new_degree = current_degree - 1

          let new_degrees = dict.insert(degrees, neighbor, new_degree)
          let new_q = case new_degree == 0 {
            True -> [neighbor, ..q]
            False -> q
          }
          #(new_q, new_degrees)
        })

      do_kahn(
        graph,
        next_queue,
        next_in_degrees,
        [head, ..acc],
        total_node_count,
      )
    }
  }
}