ex_flow_graph

defmodule ExFlow.Core.Graph do
  @moduledoc """
  Immutable graph operations wrapping libgraph with ExFlow node/edge contracts.

  ## Node Schema
  %{id: String.t(), type: atom(), label: String.t() | nil, position: %{x: number(), y: number()}, metadata: map()}

  ## Edge Schema
  %{id: String.t(), source: String.t(), source_handle: String.t(), target: String.t(), target_handle: String.t(), label: String.t() | nil, metadata: map()}
  """

  alias Graph, as: LibGraph

  @type graph_node :: %{
          id: String.t(),
          type: atom(),
          label: String.t() | nil,
          position: %{x: number(), y: number()},
          metadata: map()
        }

  @type graph_edge :: %{
          id: String.t(),
          source: String.t(),
          source_handle: String.t(),
          target: String.t(),
          target_handle: String.t(),
          label: String.t() | nil,
          metadata: map()
        }

  @type t :: LibGraph.t()

  @spec new() :: t()
  def new do
    LibGraph.new()
  end

  @spec add_node(t(), String.t(), atom(), map()) :: {:ok, t()} | {:error, term()}
  def add_node(%LibGraph{} = graph, id, type, metadata \\ %{}) when is_binary(id) do
    position = Map.get(metadata, :position, %{x: 0, y: 0})
    label = Map.get(metadata, :label)
    node_metadata = Map.get(metadata, :metadata, %{})

    node = %{
      id: id,
      type: type,
      label: label,
      position: position,
      metadata: node_metadata
    }

    case validate_node(node) do
      :ok -> {:ok, LibGraph.add_vertex(graph, id, [node])}
      error -> error
    end
  end

  @spec add_edge(t(), String.t(), String.t(), String.t(), String.t(), String.t(), map()) ::
          {:ok, t()} | {:error, term()}
  def add_edge(%LibGraph{} = graph, id, source_id, source_handle, target_id, target_handle, opts \\ %{})
      when is_binary(id) and is_binary(source_id) and is_binary(target_id) do
    edge_label = Map.get(opts, :label)
    edge_metadata = Map.get(opts, :metadata, %{})

    edge_meta = %{
      id: id,
      source: source_id,
      source_handle: source_handle,
      target: target_id,
      target_handle: target_handle,
      label: edge_label,
      metadata: edge_metadata
    }

    with :ok <- validate_edge(edge_meta),
         true <- LibGraph.has_vertex?(graph, source_id) || {:error, :source_not_found},
         true <- LibGraph.has_vertex?(graph, target_id) || {:error, :target_not_found} do
      {:ok, LibGraph.add_edge(graph, source_id, target_id, label: edge_meta)}
    else
      {:error, reason} -> {:error, reason}
      false -> {:error, :vertex_not_found}
    end
  end

  @spec update_node_position(t(), String.t(), %{x: number(), y: number()}) ::
          {:ok, t()} | {:error, term()}
  def update_node_position(%LibGraph{} = graph, id, %{x: x, y: y})
      when is_binary(id) and is_number(x) and is_number(y) do
    case LibGraph.vertex_labels(graph, id) do
      [node] when is_map(node) ->
        # Store all edges connected to this node before deleting
        in_edges = LibGraph.in_edges(graph, id)
        out_edges = LibGraph.out_edges(graph, id)

        # Update the node
        updated_node = put_in(node, [:position], %{x: x, y: y})
        graph = LibGraph.delete_vertex(graph, id)
        graph = LibGraph.add_vertex(graph, id, [updated_node])

        # Restore all edges
        graph =
          Enum.reduce(in_edges, graph, fn edge, g ->
            LibGraph.add_edge(g, edge.v1, id, label: edge.label)
          end)

        graph =
          Enum.reduce(out_edges, graph, fn edge, g ->
            LibGraph.add_edge(g, id, edge.v2, label: edge.label)
          end)

        {:ok, graph}

      _ ->
        {:error, :node_not_found}
    end
  end

  @spec update_node(t(), String.t(), map()) :: {:ok, t()} | {:error, term()}
  def update_node(%LibGraph{} = graph, id, updates) when is_binary(id) and is_map(updates) do
    case LibGraph.vertex_labels(graph, id) do
      [node] when is_map(node) ->
        # Store all edges connected to this node before deleting
        in_edges = LibGraph.in_edges(graph, id)
        out_edges = LibGraph.out_edges(graph, id)

        # Update the node with provided changes
        updated_node =
          node
          |> maybe_update(:label, updates, Map.has_key?(updates, :label))
          |> maybe_update(:metadata, updates, Map.has_key?(updates, :metadata))

        graph = LibGraph.delete_vertex(graph, id)
        graph = LibGraph.add_vertex(graph, id, [updated_node])

        # Restore all edges
        graph =
          Enum.reduce(in_edges, graph, fn edge, g ->
            LibGraph.add_edge(g, edge.v1, id, label: edge.label)
          end)

        graph =
          Enum.reduce(out_edges, graph, fn edge, g ->
            LibGraph.add_edge(g, id, edge.v2, label: edge.label)
          end)

        {:ok, graph}

      _ ->
        {:error, :node_not_found}
    end
  end

  @spec update_edge(t(), String.t(), map()) :: {:ok, t()} | {:error, term()}
  def update_edge(%LibGraph{} = graph, id, updates) when is_binary(id) and is_map(updates) do
    # Find the edge
    edge =
      get_edges(graph)
      |> Enum.find(&(&1.id == id))

    case edge do
      nil ->
        {:error, :edge_not_found}

      edge ->
        # Update edge metadata
        updated_edge =
          edge
          |> maybe_update(:label, updates, Map.has_key?(updates, :label))
          |> maybe_update(:metadata, updates, Map.has_key?(updates, :metadata))

        # Remove old edge and add updated one
        graph = LibGraph.delete_edge(graph, edge.source, edge.target)
        {:ok, LibGraph.add_edge(graph, edge.source, edge.target, label: updated_edge)}
    end
  end

  # Only update if the key is present in the updates map
  defp maybe_update(map, _key, _updates, false), do: map
  defp maybe_update(map, key, updates, true), do: Map.put(map, key, updates[key])

  @spec delete_node(t(), String.t()) :: {:ok, t()} | {:error, term()}
  def delete_node(%LibGraph{} = graph, id) when is_binary(id) do
    if LibGraph.has_vertex?(graph, id) do
      {:ok, LibGraph.delete_vertex(graph, id)}
    else
      {:error, :node_not_found}
    end
  end

  @spec delete_edge(t(), String.t()) :: {:ok, t()} | {:error, term()}
  def delete_edge(%LibGraph{} = graph, edge_id) when is_binary(edge_id) do
    case find_edge_by_id(graph, edge_id) do
      {source_id, target_id} ->
        {:ok, LibGraph.delete_edge(graph, source_id, target_id)}

      nil ->
        {:error, :edge_not_found}
    end
  end

  @spec get_node(t(), String.t()) :: {:ok, graph_node()} | {:error, term()}
  def get_node(%LibGraph{} = graph, id) when is_binary(id) do
    case LibGraph.vertex_labels(graph, id) do
      [node] when is_map(node) -> {:ok, node}
      _ -> {:error, :node_not_found}
    end
  end

  @spec get_nodes(t()) :: [graph_node()]
  def get_nodes(%LibGraph{} = graph) do
    graph
    |> LibGraph.vertices()
    |> Enum.flat_map(fn vertex_id ->
      case LibGraph.vertex_labels(graph, vertex_id) do
        [node] when is_map(node) -> [node]
        _ -> []
      end
    end)
  end

  @spec get_edges(t()) :: [graph_edge()]
  def get_edges(%LibGraph{} = graph) do
    graph
    |> LibGraph.edges()
    |> Enum.flat_map(fn edge ->
      case edge.label do
        %{id: _, source: _, target: _, source_handle: _, target_handle: _} = edge_meta ->
          [edge_meta]

        _ ->
          []
      end
    end)
  end

  @spec validate_node(graph_node()) :: :ok | {:error, term()}
  def validate_node(%{id: id, type: type, position: %{x: x, y: y}, metadata: metadata})
      when is_binary(id) and is_atom(type) and is_number(x) and is_number(y) and is_map(metadata) do
    :ok
  end

  def validate_node(_), do: {:error, :invalid_node_schema}

  @spec validate_edge(graph_edge()) :: :ok | {:error, term()}
  def validate_edge(%{id: id, source: source, source_handle: sh, target: target, target_handle: th})
      when is_binary(id) and is_binary(source) and is_binary(sh) and is_binary(target) and is_binary(th) do
    :ok
  end

  def validate_edge(_), do: {:error, :invalid_edge_schema}

  @spec to_map(t()) :: %{nodes: [graph_node()], edges: [graph_edge()]}
  def to_map(%LibGraph{} = graph) do
    %{
      nodes: get_nodes(graph),
      edges: get_edges(graph)
    }
  end

  @spec from_map(%{nodes: [graph_node()], edges: [graph_edge()]}) :: {:ok, t()} | {:error, term()}
  def from_map(%{nodes: nodes, edges: edges}) when is_list(nodes) and is_list(edges) do
    with {:ok, graph_with_nodes} <- add_nodes(new(), nodes) do
      add_edges(graph_with_nodes, edges)
    end
  end

  def from_map(_), do: {:error, :invalid_map_format}

  # Private helpers

  defp add_nodes(graph, nodes) do
    Enum.reduce_while(nodes, {:ok, graph}, fn node, {:ok, acc_graph} ->
      opts = %{
        position: node.position,
        label: Map.get(node, :label),
        metadata: node.metadata
      }

      case add_node(acc_graph, node.id, node.type, opts) do
        {:ok, new_graph} -> {:cont, {:ok, new_graph}}
        error -> {:halt, error}
      end
    end)
  end

  defp add_edges(graph, edges) do
    Enum.reduce_while(edges, {:ok, graph}, fn edge, {:ok, acc_graph} ->
      opts = %{
        label: Map.get(edge, :label),
        metadata: Map.get(edge, :metadata, %{})
      }

      case add_edge(
             acc_graph,
             edge.id,
             edge.source,
             edge.source_handle,
             edge.target,
             edge.target_handle,
             opts
           ) do
        {:ok, new_graph} -> {:cont, {:ok, new_graph}}
        error -> {:halt, error}
      end
    end)
  end

  defp find_edge_by_id(graph, edge_id) do
    graph
    |> LibGraph.edges()
    |> Enum.find_value(fn edge ->
      case edge.label do
        %{id: ^edge_id} -> {edge.v1, edge.v2}
        _ -> nil
      end
    end)
  end
end