defmodule Vela.Distributed.Ring do
  @moduledoc """
  Consistent hash ring for the Partitioned topology.

  Uses virtual nodes (vnodes) to distribute keys evenly across physical nodes.
  Each physical node gets `@vnodes_per_node` slots on the ring. To find the
  owner of a key, we hash the key and walk clockwise to the next vnode slot.

  The ring is stored as a sorted list of `{hash, node}` tuples.
  """

  @vnodes_per_node 128

  defstruct [:ring, :nodes]

  @type t :: %__MODULE__{
          ring: [{non_neg_integer(), node()}],
          nodes: [node()]
        }

  @doc "Creates a new ring from a list of nodes."
  @spec new([node()]) :: t()
  def new(nodes) when is_list(nodes) do
    ring =
      nodes
      |> Enum.flat_map(fn node ->
        for i <- 0..(@vnodes_per_node - 1) do
          hash = hash_slot(node, i)
          {hash, node}
        end
      end)
      |> Enum.sort_by(&elem(&1, 0))

    %__MODULE__{ring: ring, nodes: Enum.sort(nodes)}
  end

  @doc "Returns the node that owns the given key."
  @spec owner(t(), term()) :: node() | nil
  def owner(%__MODULE__{ring: []}, _key), do: nil

  def owner(%__MODULE__{ring: ring}, key) do
    key_hash = hash_key(key)

    # Find the first vnode with hash >= key_hash (walk clockwise)
    # If key_hash is beyond all slots, wrap around to the first slot
    case Enum.find(ring, fn {hash, _node} -> hash >= key_hash end) do
      {_hash, node} -> node
      nil -> ring |> hd() |> elem(1)
    end
  end

  @doc "Adds a node to the ring. Returns the updated ring."
  @spec add_node(t(), node()) :: t()
  def add_node(%__MODULE__{nodes: nodes}, node) do
    if node in nodes do
      new(nodes)
    else
      new([node | nodes])
    end
  end

  @doc "Removes a node from the ring. Returns the updated ring."
  @spec remove_node(t(), node()) :: t()
  def remove_node(%__MODULE__{nodes: nodes}, node) do
    new(List.delete(nodes, node))
  end

  @doc "Returns all nodes in the ring."
  @spec nodes(t()) :: [node()]
  def nodes(%__MODULE__{nodes: nodes}), do: nodes

  @doc "Returns the number of physical nodes in the ring."
  @spec size(t()) :: non_neg_integer()
  def size(%__MODULE__{nodes: nodes}), do: length(nodes)

  # ---- Internals ----

  defp hash_slot(node, index) do
    hash("#{node}:#{index}")
  end

  defp hash_key(key) do
    hash(:erlang.term_to_binary(key))
  end

  defp hash(data) do
    <<hash::unsigned-32, _::binary>> = :crypto.hash(:md5, data)
    hash
  end
end