///|
pub(all) struct GraphPath {
  nodes : Array[NodeId]
  edges : Array[EdgeId]
  tracks : Array[TrackId]
  total_distance : Double
} derive(Debug)

///|
pub fn GraphPath::empty() -> GraphPath {
  { nodes: [], edges: [], tracks: [], total_distance: 0.0 }
}

///|
pub fn GraphPath::is_valid(self : GraphPath) -> Bool {
  self.nodes.length() > 0 && self.edges.length() == self.nodes.length() - 1
}

///|
pub(all) struct StationGraph {
  nodes : Map[Int, TopologyNode]
  edges : Map[Int, TopologyEdge]
  adjacency : Map[Int, Array[EdgeId]]
  reverse_adjacency : Map[Int, Array[EdgeId]]
  tracks : Map[Int, TrackSegment]
  turnouts : Map[Int, Turnout]
  signals : Map[Int, Signal]
} derive(Debug)

///|
pub fn StationGraph::new() -> StationGraph {
  {
    nodes: Map([]),
    edges: Map([]),
    adjacency: Map([]),
    reverse_adjacency: Map([]),
    tracks: Map([]),
    turnouts: Map([]),
    signals: Map([]),
  }
}

///|
pub fn StationGraph::add_node(self : StationGraph, node : TopologyNode) -> Unit {
  self.nodes.set(node.id.id, node)
  if !self.adjacency.contains(node.id.id) {
    self.adjacency.set(node.id.id, [])
  }
  if !self.reverse_adjacency.contains(node.id.id) {
    self.reverse_adjacency.set(node.id.id, [])
  }
}

///|
pub fn StationGraph::add_edge(self : StationGraph, edge : TopologyEdge) -> Unit {
  self.edges.set(edge.id.id, edge)

  match self.adjacency.get(edge.from_node.id) {
    Some(arr) => arr.push(edge.id)
    None => self.adjacency.set(edge.from_node.id, [edge.id])
  }

  match self.reverse_adjacency.get(edge.to_node.id) {
    Some(arr) => arr.push(edge.id)
    None => self.reverse_adjacency.set(edge.to_node.id, [edge.id])
  }
}

///|
pub fn StationGraph::add_track(
  self : StationGraph,
  track : TrackSegment,
) -> Unit {
  self.tracks.set(track.id.id, track)
}

///|
pub fn StationGraph::add_turnout(
  self : StationGraph,
  turnout : Turnout,
) -> Unit {
  self.turnouts.set(turnout.id.id, turnout)
}

///|
pub fn StationGraph::add_signal(self : StationGraph, signal : Signal) -> Unit {
  self.signals.set(signal.id.id, signal)
}

///|
pub fn StationGraph::get_node(
  self : StationGraph,
  id : NodeId,
) -> TopologyNode? {
  self.nodes.get(id.id)
}

///|
pub fn StationGraph::get_edge(
  self : StationGraph,
  id : EdgeId,
) -> TopologyEdge? {
  self.edges.get(id.id)
}

///|
pub fn StationGraph::get_track(
  self : StationGraph,
  id : TrackId,
) -> TrackSegment? {
  self.tracks.get(id.id)
}

///|
pub fn StationGraph::get_turnout(
  self : StationGraph,
  id : TurnoutId,
) -> Turnout? {
  self.turnouts.get(id.id)
}

///|
pub fn StationGraph::get_signal(self : StationGraph, id : SignalId) -> Signal? {
  self.signals.get(id.id)
}

///|
pub fn StationGraph::get_out_edges(
  self : StationGraph,
  node_id : NodeId,
) -> Array[TopologyEdge] {
  match self.adjacency.get(node_id.id) {
    None => []
    Some(edge_ids) => {
      let result = []
      for edge_id in edge_ids {
        match self.edges.get(edge_id.id) {
          Some(e) => result.push(e)
          None => ()
        }
      }
      result
    }
  }
}

///|
pub fn StationGraph::get_in_edges(
  self : StationGraph,
  node_id : NodeId,
) -> Array[TopologyEdge] {
  match self.reverse_adjacency.get(node_id.id) {
    None => []
    Some(edge_ids) => {
      let result = []
      for edge_id in edge_ids {
        match self.edges.get(edge_id.id) {
          Some(e) => result.push(e)
          None => ()
        }
      }
      result
    }
  }
}

///|
pub fn StationGraph::is_turnout_path_active(
  self : StationGraph,
  from_node : NodeId,
  to_node : NodeId,
) -> Bool {
  match self.nodes.get(from_node.id) {
    Some(node) =>
      match node.associated_turnout {
        None => true
        Some(t_id) =>
          match self.turnouts.get(t_id.id) {
            None => true
            Some(t) =>
              if from_node.id == t.base_node.id {
                match t.position {
                  Normal => to_node.id == t.normal_next_node.id
                  Reverse => to_node.id == t.reverse_next_node.id
                  Moving | UnknownPosition => false
                }
              } else if to_node.id == t.base_node.id {
                if from_node.id == t.normal_next_node.id {
                  t.position == Normal
                } else if from_node.id == t.reverse_next_node.id {
                  t.position == Reverse
                } else {
                  false
                }
              } else {
                true
              }
          }
      }
    None => false
  }
}

///|
pub fn StationGraph::find_all_paths(
  self : StationGraph,
  start : NodeId,
  target : NodeId,
  max_depth : Int,
) -> Array[GraphPath] {
  let results : Array[GraphPath] = []
  let visited_nodes : Map[Int, Bool] = Map([])

  fn dfs(
    curr : NodeId,
    curr_nodes : Array[NodeId],
    curr_edges : Array[EdgeId],
    curr_tracks : Array[TrackId],
    dist : Double,
    depth : Int,
  ) -> Unit {
    if depth > max_depth {
      return
    }
    if curr.id == target.id {
      let path_nodes = curr_nodes.copy()
      let path_edges = curr_edges.copy()
      let path_tracks = curr_tracks.copy()
      results.push({
        nodes: path_nodes,
        edges: path_edges,
        tracks: path_tracks,
        total_distance: dist,
      })
      return
    }

    visited_nodes.set(curr.id, true)
    let out_edges = self.get_out_edges(curr)

    for edge in out_edges {
      let next_node = edge.to_node
      match visited_nodes.get(next_node.id) {
        Some(true) => continue
        _ => {
          curr_nodes.push(next_node)
          curr_edges.push(edge.id)
          curr_tracks.push(edge.track_id)

          dfs(
            next_node,
            curr_nodes,
            curr_edges,
            curr_tracks,
            dist + edge.weight,
            depth + 1,
          )

          let _ = curr_nodes.pop()
          let _ = curr_edges.pop()
          let _ = curr_tracks.pop()
        }
      }
    }

    visited_nodes.set(curr.id, false)
  }

  visited_nodes.set(start.id, true)
  dfs(start, [start], [], [], 0.0, 0)
  results
}

///|
pub fn StationGraph::validate_topology(self : StationGraph) -> Array[String] {
  let errors : Array[String] = []

  for _, node in self.nodes {
    let out_e = self.get_out_edges(node.id)
    let in_e = self.get_in_edges(node.id)
    if out_e.length() == 0 && in_e.length() == 0 {
      errors.push("Orphan node detected: " + node.name)
    }
  }

  for _, track in self.tracks {
    if !self.nodes.contains(track.start_node.id) {
      errors.push("Track " + track.name + " references invalid start_node")
    }
    if !self.nodes.contains(track.end_node.id) {
      errors.push("Track " + track.name + " references invalid end_node")
    }
  }

  for _, turnout in self.turnouts {
    if !self.nodes.contains(turnout.base_node.id) {
      errors.push("Turnout " + turnout.name + " references invalid base_node")
    }
    if !self.nodes.contains(turnout.normal_next_node.id) {
      errors.push(
        "Turnout " + turnout.name + " references invalid normal_next_node",
      )
    }
    if !self.nodes.contains(turnout.reverse_next_node.id) {
      errors.push(
        "Turnout " + turnout.name + " references invalid reverse_next_node",
      )
    }
  }

  errors
}