defmodule ConceptDiagram do
  @moduledoc """
  Build concept diagrams as plain Elixir data, then render them to
  [Mermaid](https://mermaid.js.org) or [Graphviz DOT](https://graphviz.org).

  A concept diagram is a set of labelled nodes connected by labelled, directed
  edges. Describe one with a small pipeline-friendly API and turn it into diagram
  source you can drop straight into Markdown (Mermaid) or render with Graphviz (DOT).

      ConceptDiagram.new(direction: :left_right)
      |> ConceptDiagram.add_edge("Idea", "Draft", "becomes")
      |> ConceptDiagram.add_edge("Draft", "Diagram", "rendered as")
      |> ConceptDiagram.to_mermaid()

  Nodes referenced by an edge are created automatically, so the snippet above is
  enough to produce a complete `flowchart LR`.

  Maintained by the team behind [Vizcept](https://vizcept.com), an AI concept
  diagram generator; this library is the code-first companion for generating
  diagram source programmatically.
  """

  @typedoc "A node id; atoms and other terms are coerced to strings."
  @type id :: String.t() | atom()

  @typedoc "Layout direction for the rendered diagram."
  @type direction :: :top_down | :bottom_up | :left_right | :right_left

  @type t :: %__MODULE__{
          nodes: [{String.t(), String.t()}],
          edges: [{String.t(), String.t(), String.t() | nil}],
          direction: direction(),
          title: String.t() | nil
        }

  defstruct nodes: [], edges: [], direction: :top_down, title: nil

  @doc """
  Create an empty diagram.

  Options:

    * `:direction` — one of `:top_down` (default), `:bottom_up`, `:left_right`,
      `:right_left`.
    * `:title` — optional title (rendered as a graph label in DOT output).
  """
  @spec new(keyword()) :: t()
  def new(opts \\ []) do
    %__MODULE__{
      direction: Keyword.get(opts, :direction, :top_down),
      title: Keyword.get(opts, :title)
    }
  end

  @doc """
  Add a node, or relabel it if the id already exists. The label defaults to the
  stringified id.
  """
  @spec add_node(t(), id(), String.t() | nil) :: t()
  def add_node(%__MODULE__{} = diagram, id, label \\ nil) do
    sid = to_string(id)
    %{diagram | nodes: put_node(diagram.nodes, sid, label || sid)}
  end

  @doc """
  Add a directed edge `from -> to` with an optional relationship label. Endpoints
  that are not yet present are added as nodes automatically.
  """
  @spec add_edge(t(), id(), id(), String.t() | nil) :: t()
  def add_edge(%__MODULE__{} = diagram, from, to, label \\ nil) do
    f = to_string(from)
    t = to_string(to)

    diagram
    |> ensure_node(f)
    |> ensure_node(t)
    |> Map.update!(:edges, &(&1 ++ [{f, t, label}]))
  end

  @doc """
  Build a diagram from a list of `{from, relation, to}` triples. `relation` may be
  `nil` for an unlabelled edge.
  """
  @spec from_triples([{id(), String.t() | nil, id()}], keyword()) :: t()
  def from_triples(triples, opts \\ []) when is_list(triples) do
    Enum.reduce(triples, new(opts), fn {from, relation, to}, diagram ->
      add_edge(diagram, from, to, relation)
    end)
  end

  @doc "Render the diagram as [Mermaid](https://mermaid.js.org) flowchart source."
  @spec to_mermaid(t()) :: String.t()
  def to_mermaid(%__MODULE__{} = diagram) do
    keys = key_map(diagram.nodes)

    nodes = for {id, label} <- diagram.nodes, do: ~s(    #{keys[id]}["#{mermaid_escape(label)}"])

    edges =
      for {f, t, label} <- diagram.edges do
        if blank?(label) do
          "    #{keys[f]} --> #{keys[t]}"
        else
          "    #{keys[f]} -->|#{mermaid_escape(label)}| #{keys[t]}"
        end
      end

    Enum.join(["flowchart #{mermaid_dir(diagram.direction)}" | nodes ++ edges], "\n")
  end

  @doc "Render the diagram as [Graphviz DOT](https://graphviz.org) source."
  @spec to_dot(t()) :: String.t()
  def to_dot(%__MODULE__{} = diagram) do
    keys = key_map(diagram.nodes)
    title = if diagram.title, do: [~s(  label="#{dot_escape(diagram.title)}";)], else: []
    nodes = for {id, label} <- diagram.nodes, do: ~s(  #{keys[id]} [label="#{dot_escape(label)}"];)

    edges =
      for {f, t, label} <- diagram.edges do
        if blank?(label) do
          "  #{keys[f]} -> #{keys[t]};"
        else
          ~s(  #{keys[f]} -> #{keys[t]} [label="#{dot_escape(label)}"];)
        end
      end

    body = ["  rankdir=#{dot_dir(diagram.direction)};"] ++ title ++ nodes ++ edges
    Enum.join(["digraph concept {" | body] ++ ["}"], "\n")
  end

  # --- internals ---

  defp put_node(nodes, id, label) do
    if List.keymember?(nodes, id, 0) do
      List.keyreplace(nodes, id, 0, {id, label})
    else
      nodes ++ [{id, label}]
    end
  end

  defp ensure_node(%__MODULE__{nodes: nodes} = diagram, id) do
    if List.keymember?(nodes, id, 0),
      do: diagram,
      else: %{diagram | nodes: nodes ++ [{id, id}]}
  end

  defp key_map(nodes) do
    nodes
    |> Enum.with_index()
    |> Map.new(fn {{id, _label}, i} -> {id, "n#{i}"} end)
  end

  defp blank?(nil), do: true
  defp blank?(""), do: true
  defp blank?(_), do: false

  defp mermaid_dir(:top_down), do: "TD"
  defp mermaid_dir(:bottom_up), do: "BT"
  defp mermaid_dir(:left_right), do: "LR"
  defp mermaid_dir(:right_left), do: "RL"

  defp dot_dir(:top_down), do: "TB"
  defp dot_dir(:bottom_up), do: "BT"
  defp dot_dir(:left_right), do: "LR"
  defp dot_dir(:right_left), do: "RL"

  defp mermaid_escape(text), do: text |> to_string() |> String.replace("\"", "#quot;")

  defp dot_escape(text) do
    text
    |> to_string()
    |> String.replace("\\", "\\\\")
    |> String.replace("\"", "\\\"")
  end
end