ex_ring_ring

defmodule ExRingRing.Infrastructure.HashRing.Static do
  @moduledoc """
  Static hash ring implementation using sorted array (tuple).

  Optimized for read-heavy workloads. Adding/removing nodes requires O(N log N) time
  but finding nodes is very fast O(log N).
  """
  @behaviour ExRingRing.Domain.Services.HashRing

  alias ExRingRing.Domain.Entities.Node
  alias ExRingRing.Domain.ValueObjects.{HashConfig, VirtualNode}
  alias ExRingRing.Infrastructure.HashRing.Base

  @type state :: %{
          vnodes: tuple(),
          base: Base.state()
        }

  @doc """
  Creates a new static hash ring.
  """
  @spec make([Node.t()], HashConfig.t()) :: state()
  def make(nodes, config) do
    base = Base.make_base_state(config)
    base = Enum.reduce(nodes, base, &Base.add_node/2)

    vnodes = build_vnodes(base)
    vnodes = List.to_tuple(vnodes)

    %{vnodes: vnodes, base: base}
  end

  @doc """
  Adds nodes to the ring.
  """
  @spec add_nodes([Node.t()], state()) :: state()
  def add_nodes(nodes, state) do
    base = Enum.reduce(nodes, state.base, &Base.add_node/2)
    vnodes = build_vnodes(base)
    %{state | vnodes: List.to_tuple(vnodes), base: base}
  end

  @doc """
  Removes nodes from the ring.
  """
  @spec remove_nodes([atom() | binary() | integer()], state()) :: state()
  def remove_nodes(keys, state) do
    base = Enum.reduce(keys, state.base, &Base.remove_node/2)
    vnodes = filter_vnodes(state.vnodes, keys)
    %{state | vnodes: vnodes, base: base}
  end

  @doc """
  Gets all nodes.
  """
  @spec get_nodes(state()) :: %{term() => Node.t()}
  def get_nodes(state), do: Base.get_nodes(state.base)

  @doc """
  Folds over nodes starting from item's position.
  """
  @spec fold(fun(), term(), term(), state()) :: term()
  def fold(fun, item, acc, state) do
    vnodes = state.vnodes
    base = state.base

    node_count = Base.get_non_phantom_node_count(base)

    if node_count == 0 do
      acc
    else
      hash_mask = Base.get_hash_mask(base)
      item_hash = Base.hash(item, base)

      partition_size = max(1, div(hash_mask + 1, tuple_size(vnodes)))
      start_pos = find_start_position(item_hash, partition_size, vnodes)

      fold_successor_nodes(node_count, start_pos, vnodes, fun, acc, MapSet.new())
    end
  end

  @doc """
  Gets non-phantom node count.
  """
  @spec get_non_phantom_node_count(state()) :: non_neg_integer()
  def get_non_phantom_node_count(state), do: Base.get_non_phantom_node_count(state.base)

  @doc """
  Hashes an item.
  """
  @spec hash(term(), state()) :: non_neg_integer()
  def hash(item, state), do: Base.hash(item, state.base)

  # Private functions

  defp build_vnodes(base) do
    nodes = Map.values(base.nodes)

    Enum.flat_map(nodes, fn node ->
      vnode_count = Base.calc_virtual_node_count(node, base)
      Base.create_virtual_nodes(node, 0, vnode_count, base)
    end)
    |> Enum.sort_by(&{&1.hash, &1.sequence})
  end

  defp filter_vnodes(vnodes, keys) do
    key_set = MapSet.new(keys)

    vnodes
    |> Tuple.to_list()
    |> Enum.reject(fn vnode -> MapSet.member?(key_set, VirtualNode.get_node_key(vnode)) end)
    |> List.to_tuple()
  end

  defp find_start_position(item_hash, partition_size, vnodes) do
    find_start_position(
      item_hash,
      partition_size,
      vnodes,
      1,
      div(item_hash, partition_size) + 1,
      tuple_size(vnodes) + 1
    )
  end

  defp find_start_position(_item_hash, _partition_size, _vnodes, pos, _, pos), do: pos

  defp find_start_position(item_hash, partition_size, vnodes, start, current, end_pos) do
    current = min(max(start, current), end_pos - 1)
    vnode = elem(vnodes, current - 1)

    if vnode.hash == item_hash do
      current
    else
      delta = item_hash - vnode.hash
      next = current + div(delta, partition_size)

      if delta > 0 do
        find_start_position(item_hash, partition_size, vnodes, current + 1, next + 1, end_pos)
      else
        find_start_position(item_hash, partition_size, vnodes, start, next - 1, current)
      end
    end
  end

  defp fold_successor_nodes(0, _pos, _vnodes, _fun, acc, _seen), do: acc

  defp fold_successor_nodes(count, pos, vnodes, fun, acc, seen) when pos > tuple_size(vnodes) do
    fold_successor_nodes(count, 1, vnodes, fun, acc, seen)
  end

  defp fold_successor_nodes(count, pos, vnodes, fun, acc, seen) do
    vnode = elem(vnodes, pos - 1)

    if MapSet.member?(seen, vnode.node) do
      fold_successor_nodes(count, pos + 1, vnodes, fun, acc, seen)
    else
      case fun.(vnode.node, acc) do
        {false, acc2} ->
          acc2

        {true, acc2} ->
          fold_successor_nodes(count - 1, pos + 1, vnodes, fun, acc2, MapSet.put(seen, vnode.node))
      end
    end
  end
end