phoenix_gen_api

defmodule PhoenixGenApi.NodeSelector do
  @moduledoc """
  Provides node selection strategies for distributed request execution.

  This module implements various strategies for selecting target nodes from a list of
  available nodes. The selection strategy determines how requests are distributed across
  nodes in a cluster.

  ## Supported Selection Strategies

  ### :random
  Selects a random node from the available nodes list. This provides simple load balancing
  with no guarantees about distribution fairness.

  ### :hash
  Uses consistent hashing based on the request ID to select a node. The same request ID
  will always map to the same node, which is useful for caching and stateful operations.

  ### {:hash, hash_key}
  Uses consistent hashing based on a specific field from the request. The `hash_key` can
  reference a field in `request.args` or a field directly on the request struct (like
  `user_id` or `device_id`). This ensures requests with the same hash_key value always
  go to the same node.

  ### :round_robin
  Distributes requests evenly across all nodes in a circular fashion. Uses an atomic
  counter for true global round-robin distribution across all processes.

  ### {:sticky, hash_key}
  Selects a node on the first request and "sticks" it to the given hash_key value
  (e.g., `user_id`). Subsequent requests with the same value go to the same node.
  The sticky mapping expires after 1 hour (configurable via `@sticky_ttl_ms`),
  after which a new node is randomly selected. If the node is removed from the
  available nodes list, a new node is selected immediately.

  ## Dynamic Node Resolution

  Instead of a static list of nodes, you can provide a module-function-args tuple that
  will be called at runtime to get the current list of nodes:

      nodes: {MyApp.NodeRegistry, :get_active_nodes, []}

  This allows for dynamic node discovery and automatic adaptation to cluster changes.

  ## Node Selection for Retry/Fallback

  The `get_nodes/2` function returns a list of nodes suitable for the retry strategy:

  - `:random` and `:hash` return a single-node list (the selected node)
  - `:round_robin` returns all nodes (for fallback to other nodes on failure)

  ## Fault Tolerance

  - Nodes are validated before selection
  - Empty node lists return `{:error, :no_nodes_available}`
  - Invalid node formats are filtered out
  - Failed node resolution is logged with details

  ## Examples

      # Random selection
      config = %FunConfig{
        nodes: ["node1@host", "node2@host", "node3@host"],
        choose_node_mode: :random
      }
      {:ok, node} = NodeSelector.get_node(config, request)

      # Hash by request ID (consistent)
      config = %FunConfig{
        nodes: ["node1@host", "node2@host"],
        choose_node_mode: :hash
      }
      {:ok, node} = NodeSelector.get_node(config, request)

      # Hash by custom field
      request = %Request{
        request_id: "req_123",
        user_id: "user_456",
        args: %{"session_id" => "sess_789"}
      }

      config = %FunConfig{
        nodes: ["node1@host", "node2@host"],
        choose_node_mode: {:hash, "user_id"}
      }
      {:ok, node} = NodeSelector.get_node(config, request)

      # Round-robin (global, atomic counter)
      config = %FunConfig{
        nodes: ["node1@host", "node2@host", "node3@host"],
        choose_node_mode: :round_robin
      }
      {:ok, node1} = NodeSelector.get_node(config, request)
      {:ok, node2} = NodeSelector.get_node(config, request)
      {:ok, node3} = NodeSelector.get_node(config, request)
      {:ok, node1_again} = NodeSelector.get_node(config, request)  # wraps around

      # Sticky node affinity
      config = %FunConfig{
        nodes: ["node1@host", "node2@host"],
        choose_node_mode: {:sticky, "user_id"}
      }
      # First request from user_123 -> randomly selects node1
      # Subsequent requests from user_123 -> always node1 (for up to 1 hour)
      {:ok, node} = NodeSelector.get_node(config, request)

  ## Notes

  - Round-robin uses an atomic counter for true global distribution (no process dictionary)
  - Hash functions use `:erlang.phash2/2` for deterministic hashing
  - Dynamic node resolution happens on every call, allowing real-time cluster updates
  - If a hash_key is not found in the request, falls back to random selection
  - Sticky mappings are stored in ETS and cleaned up periodically by `ConfigPuller` (every 1 hour)
  - Use `NodeSelector.cleanup_sticky_table/0` for manual cleanup
  - Returns `{:ok, node}` on success or `{:error, reason}` on failure
  """

  alias PhoenixGenApi.Structs.{FunConfig, Request}
  alias PhoenixGenApi.Helpers.Shared

  require Logger

  # Atomic counter reference for round-robin. Uses :atomics for lock-free increments.
  @round_robin_counter_name :phoenix_gen_api_round_robin_counter

  # ETS table for sticky node affinity.
  # Stores {sticky_key, node, timestamp} tuples.
  # sticky_key is derived from the hash_key value (e.g., user_id).
  @sticky_table_name :phoenix_gen_api_sticky_nodes
  # 1 hour TTL for sticky mappings
  @sticky_ttl_ms 3_600_000

  @doc """
  Selects a single target node based on the configuration and request.

  This function examines the `choose_node_mode` in the configuration and applies
  the appropriate node selection strategy. If the `nodes` field is a tuple, it
  will be called as a function to dynamically resolve the node list.

  ## Parameters

    - `config` - A `FunConfig` struct containing:
      - `nodes` - Either a list of node names or a `{module, function, args}` tuple
      - `choose_node_mode` - The selection strategy (`:random`, `:hash`, `{:hash, key}`, `:round_robin`)

    - `request` - A `Request` struct containing the request details

  ## Returns

    - `{:ok, node}` - The selected node
    - `{:error, reason}` - Selection failed

  ## Examples

      config = %FunConfig{
        nodes: ["node1@host", "node2@host"],
        choose_node_mode: :random
      }

      request = %Request{
        request_id: "req_123",
        request_type: "get_user",
        user_id: "user_456",
        args: %{"user_id" => "user_456"}
      }

      {:ok, node} = NodeSelector.get_node(config, request)
  """
  @spec get_node(FunConfig.t(), Request.t()) ::
          {:ok, node :: atom() | String.t()} | {:error, term()}
  def get_node(config = %FunConfig{}, request = %Request{}) do
    with {:ok, resolved_config} <- resolve_nodes(config),
         {:ok, nodes} <- validate_and_get_nodes(resolved_config) do
      select_node(nodes, resolved_config.choose_node_mode, request)
    end
  end

  @doc """
  Selects a list of target nodes based on the configuration and request.

  The returned list is ordered by preference for fallback/retry purposes:

  - `:random` - Returns a shuffled list of all nodes (selected node first)
  - `:hash` / `{:hash, key}` - Returns all nodes with the hashed node first
  - `:round_robin` - Returns all nodes starting from the round-robin position

  This is useful for the executor's fallback mechanism where if the primary
  node fails, it tries the remaining nodes in order.

  ## Parameters

    - `config` - A `FunConfig` struct
    - `request` - A `Request` struct

  ## Returns

    - `{:ok, [node, ...]}` - Ordered list of nodes (primary first)
    - `{:error, reason}` - Selection failed

  ## Examples

      config = %FunConfig{
        nodes: ["node1@host", "node2@host", "node3@host"],
        choose_node_mode: :random
      }

      {:ok, [primary | fallbacks]} = NodeSelector.get_nodes(config, request)
  """
  @spec get_nodes(FunConfig.t(), Request.t()) :: {:ok, [atom() | String.t()]} | {:error, term()}
  def get_nodes(config = %FunConfig{}, request = %Request{}) do
    with {:ok, resolved_config} <- resolve_nodes(config),
         {:ok, nodes} <- validate_and_get_nodes(resolved_config) do
      select_nodes_ordered(nodes, resolved_config.choose_node_mode, request)
    end
  end

  @doc """
  Resolves dynamic node configuration to a concrete node list.

  If `nodes` is an MFA tuple `{module, function, args}`, calls the function
  to get the node list at runtime. If `nodes` is already a list, returns
  the config unchanged. If `nodes` is `:local`, returns the config unchanged.

  ## Parameters

    - `config` - A `FunConfig` struct

  ## Returns

    - `{:ok, %FunConfig{}}` - Config with resolved nodes
    - `{:error, reason}` - Resolution failed
  """
  @spec resolve_nodes(FunConfig.t()) :: {:ok, FunConfig.t()} | {:error, term()}
  def resolve_nodes(config = %FunConfig{nodes: :local}), do: {:ok, config}
  def resolve_nodes(config = %FunConfig{nodes: nodes}) when is_list(nodes), do: {:ok, config}

  def resolve_nodes(config = %FunConfig{nodes: {m, f, a}}) do
    case resolve_dynamic_nodes(m, f, a) do
      {:ok, nodes} ->
        {:ok, %{config | nodes: nodes}}

      {:error, reason} ->
        Logger.error(
          "[NodeSelector] resolve_nodes: failed to resolve dynamic nodes via MFA, reason: #{inspect(reason)}"
        )

        {:error, {:dynamic_node_resolution_failed, reason}}
    end
  end

  def resolve_nodes(%FunConfig{nodes: other}) do
    Logger.error("[NodeSelector] resolve_nodes: invalid nodes configuration: #{inspect(other)}")

    {:error, {:invalid_nodes_configuration, other}}
  end

  @doc """
  Resolves nodes and returns the raw node list regardless of configuration type.

  Unlike `resolve_nodes/1` which returns the full config, this returns just
  the list of nodes. Useful for getting all available nodes for retry strategies.

  ## Parameters

    - `config` - A `FunConfig` struct

  ## Returns

    - `{:ok, [node]}` - List of resolved nodes
    - `{:error, reason}` - Resolution failed
  """
  @spec resolve_nodes_list(FunConfig.t()) :: {:ok, [atom() | String.t()]} | {:error, term()}
  def resolve_nodes_list(%FunConfig{nodes: :local}) do
    {:ok, [node()]}
  end

  def resolve_nodes_list(config = %FunConfig{}) do
    case resolve_nodes(config) do
      {:ok, resolved} -> {:ok, Shared.validate_nodes(resolved.nodes)}
      {:error, reason} -> {:error, reason}
    end
  end

  @doc """
  Validates that the `choose_node_mode` is a recognized strategy.

  ## Returns

    - `true` if the mode is valid
    - `false` otherwise
  """
  @spec choose_node_valid?(FunConfig.t()) :: boolean()
  def choose_node_valid?(%FunConfig{choose_node_mode: mode}) do
    case mode do
      :random -> true
      :hash -> true
      {:hash, _} -> true
      :round_robin -> true
      {:sticky, _} -> true
      _ -> false
    end
  end

  @doc """
  Calculates a retry backoff delay based on the attempt number.

  Uses exponential backoff with jitter to prevent thundering herd problems.

  ## Parameters

    - `attempt` - The current attempt number (1-based)
    - `opts` - Options keyword list:
      - `:base_ms` - Base delay in milliseconds (default: 100)
      - `:max_ms` - Maximum delay in milliseconds (default: 5000)
      - `:jitter` - Whether to add random jitter (default: true)

  ## Returns

    - Delay in milliseconds before the next retry

  ## Examples

      iex> NodeSelector.calculate_backoff(1)
      # Returns ~100ms (with jitter)

      iex> NodeSelector.calculate_backoff(3)
      # Returns ~400ms (with jitter)

      iex> NodeSelector.calculate_backoff(5, max_ms: 2000)
      # Returns capped at ~2000ms (with jitter)
  """
  @spec calculate_backoff(pos_integer(), keyword()) :: non_neg_integer()
  def calculate_backoff(attempt, opts \\ []) when is_integer(attempt) and attempt > 0 do
    base_ms = Keyword.get(opts, :base_ms, 100)
    max_ms = Keyword.get(opts, :max_ms, 5_000)
    jitter? = Keyword.get(opts, :jitter, true)

    # Exponential backoff: base * 2^(attempt-1)
    delay = (base_ms * :math.pow(2, attempt - 1)) |> trunc()
    delay = min(delay, max_ms)

    if jitter? do
      # Add random jitter: 0.5x to 1.5x the delay
      jitter_factor = 0.5 + :rand.uniform()
      trunc(delay * jitter_factor)
    else
      delay
    end
  end

  @doc """
  Resets the round-robin counter.

  This is primarily useful for testing. In production, the counter
  should be allowed to increment naturally.
  """
  @spec reset_round_robin() :: :ok
  def reset_round_robin do
    case :ets.whereis(@round_robin_counter_name) do
      :undefined ->
        :ok

      _table ->
        :ets.insert(@round_robin_counter_name, {:counter, 0})
        :ok
    end
  end

  # --- Private Functions ---

  defp resolve_dynamic_nodes(m, f, a) when is_atom(m) and is_atom(f) and is_list(a) do
    try do
      case apply(m, f, a) do
        nodes when is_list(nodes) ->
          {:ok, nodes}

        other ->
          {:error, {:invalid_return_type, other}}
      end
    rescue
      error ->
        {:error, {:exception, Exception.message(error)}}
    catch
      kind, value ->
        {:error, {kind, value}}
    end
  end

  defp resolve_dynamic_nodes(_, _, _) do
    {:error, :invalid_mfa_format}
  end

  defp validate_and_get_nodes(%FunConfig{nodes: :local}) do
    {:ok, [node()]}
  end

  defp validate_and_get_nodes(%FunConfig{nodes: nodes}) do
    validated = Shared.validate_nodes(nodes)

    if validated == [] do
      Logger.error("[NodeSelector] no valid nodes available after filtering")
      {:error, :no_nodes_available}
    else
      {:ok, validated}
    end
  end

  # Select a single node (backward compatible with get_node/2)
  defp select_node(nodes, mode, request) do
    case mode do
      :random ->
        {:ok, Enum.random(nodes)}

      :hash ->
        {:ok, hash_node(request, nodes)}

      {:hash, hash_key} ->
        hash_node_with_fallback(request, nodes, hash_key)

      :round_robin ->
        {:ok, round_robin_node(nodes)}

      {:sticky, hash_key} ->
        sticky_node(request, nodes, hash_key)

      _ ->
        Logger.error("[NodeSelector] invalid choose_node_mode: #{inspect(mode)}")

        {:error, {:invalid_choose_node_mode, mode}}
    end
  end

  # Select an ordered list of nodes (primary first, then fallbacks)
  defp select_nodes_ordered(nodes, mode, request) do
    case mode do
      :random ->
        # Shuffle nodes, putting a random one first
        primary = Enum.random(nodes)
        fallbacks = List.delete(nodes, primary) |> Enum.shuffle()
        {:ok, [primary | fallbacks]}

      :hash ->
        primary = hash_node(request, nodes)
        fallbacks = List.delete(nodes, primary)
        {:ok, [primary | fallbacks]}

      {:hash, hash_key} ->
        # hash_node_with_fallback always returns {:ok, node} (falling back to random on miss)
        {:ok, primary} = hash_node_with_fallback(request, nodes, hash_key)
        fallbacks = List.delete(nodes, primary)
        {:ok, [primary | fallbacks]}

      :round_robin ->
        # Return all nodes starting from the round-robin position
        primary_idx = get_round_robin_index(length(nodes))
        {before, after_nodes} = Enum.split(nodes, primary_idx)
        {:ok, after_nodes ++ before}

      {:sticky, hash_key} ->
        case sticky_node(request, nodes, hash_key) do
          {:ok, primary} ->
            fallbacks = List.delete(nodes, primary)
            {:ok, [primary | fallbacks]}

          error ->
            error
        end

      _ ->
        Logger.error("[NodeSelector] invalid choose_node_mode: #{inspect(mode)}")

        {:error, {:invalid_choose_node_mode, mode}}
    end
  end

  defp hash_node(request, nodes) do
    hash_order = :erlang.phash2(request.request_id, length(nodes))
    Enum.at(nodes, hash_order)
  end

  defp hash_node_with_fallback(request, nodes, hash_key) do
    value =
      Map.get(request.args, hash_key) ||
        case request do
          %{^hash_key => v} when not is_nil(v) -> v
          _ -> nil
        end

    case value do
      nil ->
        Logger.warning(
          "[NodeSelector] hash key #{inspect(hash_key)} not found in request args, falling back to random selection"
        )

        {:ok, Enum.random(nodes)}

      val ->
        hash_order = :erlang.phash2(val, length(nodes))
        {:ok, Enum.at(nodes, hash_order)}
    end
  end

  # Round-robin using an atomic counter stored in ETS.
  # This provides true global round-robin across all processes,
  # unlike the previous process dictionary approach.
  defp round_robin_node(nodes) do
    nodes_length = length(nodes)
    idx = get_round_robin_index(nodes_length)
    Enum.at(nodes, idx)
  end

  defp get_round_robin_index(nodes_length) when nodes_length <= 1, do: 0

  defp get_round_robin_index(nodes_length) do
    ensure_round_robin_counter()

    # Use :ets.update_counter for atomic lock-free increment.
    # :ets.update_counter/4 uses the same underlying mechanism as :atomics
    # (lock-free atomic operations on ETS entries) and scales linearly with cores.
    # This avoids the single-writer bottleneck of a GenServer-based counter.
    counter =
      :ets.update_counter(
        @round_robin_counter_name,
        :counter,
        {2, 1},
        {@round_robin_counter_name, 0}
      )

    # counter is the value AFTER the increment, so use (counter - 1) for 0-based indexing
    rem(counter - 1, nodes_length)
  end

  # Idempotent initialization of the round-robin counter ETS table.
  # Uses :ets.update_counter for lock-free atomic increments (same underlying
  # mechanism as :atomics, but without the ETS table ownership issue).
  defp ensure_round_robin_counter do
    case :ets.whereis(@round_robin_counter_name) do
      :undefined ->
        try do
          :ets.new(@round_robin_counter_name, [
            :named_table,
            :public,
            :set,
            write_concurrency: true,
            read_concurrency: true
          ])

          :ets.insert(@round_robin_counter_name, {:counter, 0})
        rescue
          ArgumentError ->
            # Table was created by another process concurrently; that's fine
            :ok
        end

      _ ->
        :ok
    end
  end

  # --- Sticky Node Affinity ---

  # Gets or creates a sticky node assignment for the given hash_key value.
  # Returns {:ok, node} or {:error, reason}.
  defp sticky_node(request, nodes, hash_key) do
    value = get_sticky_value(request, hash_key)

    if is_nil(value) do
      Logger.warning(
        "[NodeSelector] sticky hash key #{inspect(hash_key)} not found in request, falling back to random selection"
      )

      select_and_store_sticky(nodes, hash_key)
    else
      ensure_sticky_table()
      handle_sticky_lookup(nodes, hash_key)
    end
  end

  defp handle_sticky_lookup(nodes, sticky_key) do
    case :ets.lookup(@sticky_table_name, sticky_key) do
      [{^sticky_key, node, _ts}] ->
        handle_existing_sticky(nodes, node, sticky_key)

      _ ->
        # No sticky assignment or node no longer in list, select new one
        select_and_store_sticky(nodes, sticky_key)
    end
  end

  defp handle_existing_sticky(nodes, node, sticky_key) do
    if Enum.member?(nodes, node) do
      if sticky_valid?(sticky_key) do
        {:ok, node}
      else
        # TTL expired, re-select
        select_and_store_sticky(nodes, sticky_key)
      end
    else
      # Node no longer in list, select new one
      select_and_store_sticky(nodes, sticky_key)
    end
  end

  defp get_sticky_value(request, hash_key) do
    Map.get(request.args, hash_key) ||
      case request do
        %{^hash_key => v} when not is_nil(v) -> v
        _ -> nil
      end
  end

  defp ensure_sticky_table do
    case :ets.whereis(@sticky_table_name) do
      :undefined ->
        try do
          :ets.new(@sticky_table_name, [
            :named_table,
            :public,
            :set,
            write_concurrency: true
          ])
        rescue
          ArgumentError ->
            # Table was created by another process concurrently
            :ok
        end

      _ ->
        :ok
    end
  end

  defp sticky_valid?(sticky_key) do
    case :ets.lookup(@sticky_table_name, sticky_key) do
      [{^sticky_key, _node, ts}] ->
        now_ms = System.system_time(:millisecond)
        now_ms - ts < @sticky_ttl_ms

      _ ->
        false
    end
  end

  defp select_and_store_sticky(nodes, sticky_key) do
    ensure_sticky_table()
    node = Enum.random(nodes)
    ts = System.system_time(:millisecond)
    :ets.insert(@sticky_table_name, {sticky_key, node, ts})
    {:ok, node}
  end

  @doc """
  Cleans up expired sticky mappings.

  This function should be called periodically (e.g., every hour) to remove
  stale sticky node assignments that have exceeded their TTL.

  ## Returns

  - `:ok` - Cleanup completed
  """
  def cleanup_sticky_table do
    now_ms = System.system_time(:millisecond)

    :ets.foldl(@sticky_table_name, [], fn {key, node, ts}, acc ->
      if ts < now_ms - @sticky_ttl_ms do
        [{key, node, ts} | acc]
      else
        acc
      end
    end)
    |> Enum.each(fn {key, node, ts} ->
      :ets.delete_object(@sticky_table_name, {key, node, ts})
    end)

    :ok
  end
end