ex_ai

defmodule ExAI.Graph do
  @moduledoc """
  Public graph builder and runtime API.

  `ExAI.Graph` is the foundation runtime used by higher-level systems in ExAI.

  ## Minimal example

      graph =
        ExAI.Graph.new(id: :hello_graph, state: %{value: 1})
        |> ExAI.Graph.step(:inc, fn state -> {:ok, %{value: state.value + 1}} end)
        |> ExAI.Graph.edge(:start, :inc)
        |> ExAI.Graph.edge(:inc, :end)

      {:ok, context} = ExAI.Graph.run(graph)
      context.result
      #=> %{value: 2}

  ## Useful helpers

    * `run/3` and `run!/3` for full execution
    * `iter/3` and `iter_next/1` for step-by-step execution
    * `mermaid/2` for graph visualization
  """

  alias ExAI.Graph.Builder
  alias ExAI.Graph.Definition
  alias ExAI.Types

  @type t :: Definition.t()

  @typedoc "Step callback input shape."
  @type step_callback ::
          (term() -> ExAI.Graph.Step.result())
          | (term(), Types.NodeContext.t() -> ExAI.Graph.Step.result())
          | Definition.node_callback()

  @typedoc "Decision callback that returns a route/branch identifier."
  @type decision_callback ::
          (term() -> Types.node_id())
          | (term(), Types.NodeContext.t() -> Types.node_id())
          | Definition.node_callback()

  @doc """
  Starts a new graph definition.
  """
  @spec new(keyword() | map()) :: t()
  def new(opts \\ []), do: Builder.new(opts)

  @doc """
  Registers a step node.
  """
  @spec step(t(), Types.node_id(), step_callback() | keyword() | map(), keyword() | map()) :: t()
  def step(%Definition{} = graph, id, callback_or_opts \\ [], opts \\ []) do
    {callback, normalized_opts} = split_callback_and_opts(callback_or_opts, opts)
    Builder.add_step(graph, id, callback, normalized_opts)
  end

  @doc """
  Registers a decision node.
  """
  @spec decision(
          t(),
          Types.node_id(),
          decision_callback() | keyword() | map(),
          keyword() | map()
        ) :: t()
  def decision(%Definition{} = graph, id, callback_or_opts \\ [], opts \\ []) do
    {callback, normalized_opts} = split_callback_and_opts(callback_or_opts, opts)

    if is_nil(callback) do
      raise ArgumentError, "decision nodes require a callback"
    end

    Builder.add_decision(graph, id, callback, normalized_opts)
  end

  @doc """
  Registers a directed edge from one node to another.
  """
  @spec edge(t(), Types.node_id(), Types.node_id(), keyword() | map()) :: t()
  def edge(%Definition{} = graph, from, to, opts \\ []) do
    Builder.add_edge(graph, from, to, opts)
  end

  @doc """
  Merges graph-level metadata into the definition.
  """
  @spec metadata(t(), keyword() | map()) :: t()
  def metadata(%Definition{} = graph, metadata) do
    Builder.merge_metadata(graph, metadata)
  end

  @doc """
  Sets a single graph-level metadata key.
  """
  @spec put_metadata(t(), atom() | String.t(), term()) :: t()
  def put_metadata(%Definition{} = graph, key, value) do
    Builder.put_metadata(graph, key, value)
  end

  @doc """
  Validates a graph definition.
  """
  @spec validate(t()) :: :ok | {:error, ExAI.Error.t()}
  def validate(%Definition{} = graph) do
    ExAI.Graph.Validator.validate(graph)
  end

  @doc """
  Validates a graph definition and raises on failure.
  """
  @spec validate!(t()) :: t()
  def validate!(%Definition{} = graph) do
    :ok = ExAI.Graph.Validator.validate(graph)
    graph
  end

  @doc """
  Runs a validated graph locally.

  ## Options

    * `:max_steps` - maximum executed nodes before timing out (default `10_000`)
    * `:timeout_ms` - wall-clock timeout for the full run
    * `:cancel_hook` - function (`arity 0` or `arity 1`) checked before each step;
      returning `true` or `{:cancel, reason}` stops the run
  """
  @spec run(t(), term(), keyword()) ::
          {:ok, ExAI.Graph.Context.t()}
          | {:error, ExAI.Error.t(), ExAI.Graph.Context.t()}
  def run(%Definition{} = graph, input \\ nil, opts \\ []) do
    ExAI.Graph.Runner.run(graph, input, opts)
  end

  @doc """
  Runs a graph locally and raises on failure.
  """
  @spec run!(t(), term(), keyword()) :: ExAI.Graph.Context.t()
  def run!(%Definition{} = graph, input \\ nil, opts \\ []) do
    ExAI.Graph.Runner.run!(graph, input, opts)
  end

  @doc """
  Renders a graph definition as Mermaid `flowchart` text.
  """
  @spec mermaid(t(), keyword() | map()) :: String.t()
  def mermaid(%Definition{} = graph, opts \\ []) do
    ExAI.Graph.Mermaid.render(graph, opts)
  end

  @doc """
  Creates a step iterator for a graph run.
  """
  @spec iter(t(), term(), keyword()) ::
          {:ok, ExAI.Graph.Iter.t()} | {:error, ExAI.Error.t(), ExAI.Graph.Iter.t()}
  def iter(%Definition{} = graph, input \\ nil, opts \\ []) do
    ExAI.Graph.Iter.new(graph, input, opts)
  end

  @doc """
  Advances a graph iterator by one node execution.
  """
  @spec iter_next(ExAI.Graph.Iter.t()) ::
          {:ok, ExAI.Graph.Iter.t()}
          | {:done, ExAI.Graph.Iter.t()}
          | {:error, ExAI.Error.t(), ExAI.Graph.Iter.t()}
  def iter_next(%ExAI.Graph.Iter{} = iterator) do
    ExAI.Graph.Iter.next(iterator)
  end

  @doc """
  Returns the node registered with the given id, if present.
  """
  @spec node(t(), Types.node_id()) :: Definition.Node.t() | nil
  def node(%Definition{} = graph, id), do: Definition.node(graph, id)

  @doc """
  Returns all registered nodes.
  """
  @spec nodes(t()) :: %{required(Types.node_id()) => Definition.Node.t()}
  def nodes(%Definition{} = graph), do: Definition.all_nodes(graph)

  @doc """
  Returns an inspectable snapshot map of the graph definition.
  """
  @spec describe(t()) :: map()
  def describe(%Definition{} = graph) do
    %{
      id: graph.id,
      state: graph.state,
      metadata: graph.metadata,
      steps:
        graph.steps
        |> sort_by_id()
        |> Enum.map(fn {id, node} ->
          %{id: id, executor: node.executor, placement: node.placement, metadata: node.metadata}
        end),
      decisions:
        graph.decisions
        |> sort_by_id()
        |> Enum.map(fn {id, node} ->
          %{id: id, executor: node.executor, placement: node.placement, metadata: node.metadata}
        end),
      edges:
        Enum.map(graph.edges, fn edge ->
          %{from: edge.from, to: edge.to, metadata: edge.metadata}
        end),
      registrations: graph.registrations,
      counts: %{
        steps: map_size(graph.steps),
        decisions: map_size(graph.decisions),
        edges: length(graph.edges)
      }
    }
  end

  @spec split_callback_and_opts(term(), keyword() | map()) :: {term() | nil, keyword() | map()}
  defp split_callback_and_opts(callback_or_opts, opts)

  defp split_callback_and_opts(callback_or_opts, opts)
       when (is_list(callback_or_opts) or is_map(callback_or_opts)) and opts != [] and opts != %{} do
    raise ArgumentError,
          "pass node options as either the 3rd argument OR the 4th argument, not both"
  end

  defp split_callback_and_opts(callback, opts) when opts != [] and opts != %{} do
    if callback_term?(callback) do
      {callback, opts}
    else
      raise ArgumentError, "invalid callback #{inspect(callback)}"
    end
  end

  defp split_callback_and_opts(callback_or_opts, _opts)
       when is_list(callback_or_opts) or is_map(callback_or_opts) do
    {nil, callback_or_opts}
  end

  defp split_callback_and_opts(callback, _opts) do
    if callback_term?(callback) do
      {callback, []}
    else
      raise ArgumentError, "invalid callback #{inspect(callback)}"
    end
  end

  @spec sort_by_id(map()) :: list()
  defp sort_by_id(map) do
    Enum.sort_by(map, fn {id, _node} -> to_string(id) end)
  end

  @spec callback_term?(term()) :: boolean()
  defp callback_term?(callback) do
    is_function(callback, 1) or
      is_function(callback, 2) or
      match?(
        {module, function, args} when is_atom(module) and is_atom(function) and is_list(args),
        callback
      ) or
      is_nil(callback)
  end
end