baileys_ex

defmodule BaileysEx.Auth.KeyStore do
  @moduledoc """
  Persistence-backed transactional Signal key store.

  This module wraps an auth persistence backend with the same `get/3`, `set/2`,
  and `transaction/3` shape used by the runtime `Signal.Store` contract. Reads
  go through ETS, transaction work is cached on an explicit transaction handle, and commit
  failures roll back to the previous persisted snapshot before surfacing an
  error to the caller. When a persistence backend exports context-aware
  callbacks such as `load_keys/3` or `save_keys/4`, the store passes the
  configured `:persistence_context` as the first argument; otherwise it falls
  back to the behaviour's context-free callbacks.
  """

  use GenServer

  @behaviour BaileysEx.Signal.Store

  alias BaileysEx.Auth.NativeFilePersistence
  alias BaileysEx.Signal.Store.LockManager
  alias BaileysEx.Signal.Store.TransactionBuffer

  @missing :"$missing"

  defmodule OperationError do
    @moduledoc false

    defexception [:action, :reason]

    @impl true
    def message(%__MODULE__{action: action, reason: reason}) do
      "auth key store #{action} failed: #{inspect(reason)}"
    end
  end

  defmodule Ref do
    @moduledoc """
    Store reference returned by `wrap/1` and passed into the KeyStore operations.
    """

    @enforce_keys [:pid, :table]
    defstruct [:pid, :table]

    @typedoc "Opaque auth key store reference."
    @type t :: %__MODULE__{pid: pid(), table: :ets.tid()}
  end

  defmodule TxRef do
    @moduledoc """
    Internal transaction-scoped KeyStore handle.

    Callers receive this only inside `transaction/3` callbacks. It carries the
    explicit transaction-local cache and mutation buffer used by the built-in
    persistence-backed Signal store implementation.

    This module exists to make the transaction contract and generated
    documentation accurate for advanced custom store implementers. Application
    code should not construct or persist these structs directly.
    """

    @enforce_keys [:pid, :table, :tx_table]
    defstruct [:pid, :table, :tx_table]

    @typedoc "Internal transaction-scoped auth key store reference."
    @type t :: %__MODULE__{pid: pid(), table: :ets.tid(), tx_table: :ets.tid()}
  end

  @type state :: %{
          table: :ets.tid(),
          persistence_module: module(),
          persistence_context: term(),
          locks: map(),
          monitor_keys: map(),
          known_ids: map(),
          max_commit_retries: pos_integer(),
          delay_between_tries_ms: non_neg_integer()
        }

  @doc """
  Starts the transactional key store linked to the current process.

  Options:

  - `:persistence_module` - module implementing `BaileysEx.Auth.Persistence`
    (defaults to `BaileysEx.Auth.NativeFilePersistence`)
  - `:persistence_context` - backend-specific context passed to the built-in
    context-aware persistence callbacks when exported
  """
  @impl true
  @spec start_link(keyword()) :: GenServer.on_start()
  def start_link(opts \\ []) do
    GenServer.start_link(__MODULE__, opts, name: Keyword.get(opts, :name))
  end

  @doc """
  Creates a read-only query ref struct to pass directly into reads.
  """
  @impl true
  @spec wrap(pid()) :: Ref.t()
  def wrap(pid) when is_pid(pid) do
    %Ref{pid: pid, table: GenServer.call(pid, :table)}
  end

  @doc """
  Fetches an array of identifiers for a given data type.
  """
  @impl true
  @spec get(Ref.t() | TxRef.t(), BaileysEx.Signal.Store.data_type(), [String.t()]) ::
          BaileysEx.Signal.Store.data_entries()
  def get(%Ref{} = ref, type, ids) when is_list(ids) do
    {entries, missing_ids} = read_cached_entries(ref.table, type, ids)
    merge_fetched_missing(entries, ref, type, missing_ids)
  end

  def get(%TxRef{} = ref, type, ids) when is_list(ids) do
    read_entries_in_transaction(ref, type, ids)
  end

  @doc """
  Sets arbitrary mutations into the persistence backend.
  """
  @impl true
  @spec set(Ref.t() | TxRef.t(), BaileysEx.Signal.Store.data_set()) :: :ok
  def set(%Ref{} = ref, data) when is_map(data) do
    case GenServer.call(ref.pid, {:set, data}, :infinity) do
      :ok -> :ok
      {:error, reason} -> raise OperationError, action: :set, reason: reason
    end
  end

  def set(%TxRef{} = ref, data) when is_map(data) do
    :ok = merge_transaction_data(ref, data)
    :ok
  end

  @doc """
  Clears all keys from persistence.
  """
  @impl true
  @spec clear(Ref.t() | TxRef.t()) :: :ok
  def clear(%Ref{} = ref) do
    case GenServer.call(ref.pid, :clear, :infinity) do
      :ok -> :ok
      {:error, reason} -> raise OperationError, action: :clear, reason: reason
    end
  end

  def clear(%TxRef{} = ref) do
    :ok = TransactionBuffer.clear(ref.tx_table)
    :ok
  end

  @doc """
  Acquires an exclusive lock tied to `key` before running the `fun`.
  Errors safely roll back changes if commit fails.
  """
  @impl true
  @spec transaction(Ref.t() | TxRef.t(), String.t(), (TxRef.t() -> result)) :: result
        when result: var
  def transaction(%Ref{} = ref, key, fun) when is_binary(key) and is_function(fun, 1) do
    :ok = GenServer.call(ref.pid, {:lock, key, self()}, :infinity)
    tx_table = TransactionBuffer.new()
    tx_ref = %TxRef{pid: ref.pid, table: ref.table, tx_table: tx_table}

    try do
      result = fun.(tx_ref)

      case GenServer.call(
             ref.pid,
             {:commit_tx, TransactionBuffer.cleared?(tx_table),
              TransactionBuffer.mutation_data(tx_table)},
             :infinity
           ) do
        :ok -> result
        {:error, reason} -> raise OperationError, action: :transaction, reason: reason
      end
    after
      TransactionBuffer.delete(tx_table)
      :ok = GenServer.call(ref.pid, {:unlock, key, self()}, :infinity)
    end
  end

  def transaction(%TxRef{} = ref, _key, fun) when is_function(fun, 1), do: fun.(ref)

  @doc """
  Returns true if the current process is in an active transaction context.
  """
  @impl true
  @spec in_transaction?(Ref.t() | TxRef.t()) :: boolean()
  def in_transaction?(%Ref{}), do: false
  def in_transaction?(%TxRef{}), do: true

  @impl true
  def init(opts) do
    table = :ets.new(__MODULE__, [:set, :protected, read_concurrency: true])

    {:ok,
     %{
       table: table,
       persistence_module: Keyword.get(opts, :persistence_module, NativeFilePersistence),
       persistence_context: Keyword.get(opts, :persistence_context),
       locks: %{},
       monitor_keys: %{},
       known_ids: %{},
       max_commit_retries: Keyword.get(opts, :max_commit_retries, 10),
       delay_between_tries_ms: Keyword.get(opts, :delay_between_tries_ms, 3_000)
     }}
  end

  @impl true
  def handle_call(:table, _from, state), do: {:reply, state.table, state}

  def handle_call({:fetch_missing, type, ids}, _from, state) do
    case fetch_missing(state, type, ids) do
      {:ok, fetched, state} -> {:reply, {:ok, fetched}, state}
      {:error, reason, state} -> {:reply, {:error, reason}, state}
    end
  end

  def handle_call({:set, data}, _from, state) do
    case commit_with_retry(state, data) do
      {:ok, state} -> {:reply, :ok, state}
      {:error, reason, state} -> {:reply, {:error, reason}, state}
    end
  end

  def handle_call({:commit_tx, clear?, data}, _from, state) do
    case commit_with_retry(state, data, clear?) do
      {:ok, state} -> {:reply, :ok, state}
      {:error, reason, state} -> {:reply, {:error, reason}, state}
    end
  end

  def handle_call(:clear, _from, state) do
    case clear_persisted_entries(state) do
      {:ok, state} -> {:reply, :ok, state}
      {:error, reason, state} -> {:reply, {:error, reason}, state}
    end
  end

  def handle_call({:lock, key, owner}, from, state) do
    case LockManager.acquire(state, key, from, owner) do
      {:acquired, updated_state} ->
        {:reply, :ok, updated_state}

      {:queued, updated_state} ->
        {:noreply, updated_state}
    end
  end

  def handle_call({:unlock, key, owner}, _from, state) do
    {:reply, :ok, LockManager.release(state, key, owner)}
  end

  @impl true
  def handle_info({:DOWN, monitor_ref, :process, owner, _reason}, state) do
    {:noreply, LockManager.handle_owner_down(state, monitor_ref, owner)}
  end

  defp read_entries_in_transaction(%TxRef{} = ref, type, ids) do
    {_entries, missing_ids} = TransactionBuffer.cached_entries(ref.tx_table, type, ids)

    if missing_ids != [] do
      fetched = fetch_transaction_missing(ref, type, missing_ids)
      :ok = TransactionBuffer.cache_fetched(ref.tx_table, type, missing_ids, fetched)
    end

    ref.tx_table
    |> TransactionBuffer.cached_entries(type, ids)
    |> elem(0)
  end

  defp fetch_transaction_missing(%TxRef{} = ref, type, missing_ids) do
    if TransactionBuffer.cleared?(ref.tx_table) do
      %{}
    else
      case GenServer.call(ref.pid, {:fetch_missing, type, missing_ids}, :infinity) do
        {:ok, result} -> result
        {:error, reason} -> raise OperationError, action: :get, reason: reason
      end
    end
  end

  defp merge_transaction_data(ref, data) do
    Enum.reduce(data, :ok, fn {type, entries}, :ok ->
      if type == :"pre-key" do
        merge_transaction_prekeys(ref, entries)
      else
        merge_transaction_entries(ref, type, entries)
      end
    end)
  end

  defp merge_transaction_entries(%TxRef{} = ref, type, entries) do
    TransactionBuffer.put_entries(ref.tx_table, %{type => entries})
  end

  defp merge_transaction_prekeys(%TxRef{} = ref, entries) do
    Enum.reduce(entries, :ok, fn
      {id, nil}, :ok ->
        case get(ref, :"pre-key", [id]) do
          %{^id => _existing} -> TransactionBuffer.put_entry(ref.tx_table, :"pre-key", id, nil)
          %{} -> :ok
        end

      {id, value}, :ok ->
        TransactionBuffer.put_entry(ref.tx_table, :"pre-key", id, value)
    end)
  end

  defp read_cached_entries(table, type, ids) do
    Enum.reduce(ids, {%{}, []}, fn id, {entries, missing_ids} ->
      case lookup_cache(table, type, id) do
        {:ok, value} -> {Map.put(entries, id, value), missing_ids}
        :cached_missing -> {entries, missing_ids}
        :miss -> {entries, [id | missing_ids]}
      end
    end)
    |> then(fn {entries, missing_ids} -> {entries, Enum.reverse(missing_ids)} end)
  end

  defp lookup_cache(table, type, id) do
    case :ets.lookup(table, {type, id}) do
      [{{^type, ^id}, @missing}] -> :cached_missing
      [{{^type, ^id}, value}] -> {:ok, value}
      [] -> :miss
    end
  end

  defp fetch_missing(state, _type, []), do: {:ok, %{}, state}

  defp fetch_missing(state, type, ids) do
    Enum.reduce_while(ids, {:ok, %{}, state}, fn id, {:ok, fetched, acc_state} ->
      case fetch_missing_id(acc_state, type, id) do
        {:ok, nil, next_state} ->
          {:cont, {:ok, fetched, next_state}}

        {:ok, value, next_state} ->
          {:cont, {:ok, Map.put(fetched, id, value), next_state}}

        {:error, reason, next_state} ->
          {:halt, {:error, reason, next_state}}
      end
    end)
  end

  defp fetch_missing_id(state, type, id) do
    case lookup_cache(state.table, type, id) do
      {:ok, value} -> {:ok, value, state}
      :cached_missing -> {:ok, nil, state}
      :miss -> load_and_cache(state, type, id)
    end
  end

  defp load_and_cache(state, type, id) do
    case persistence_load(state, type, id) do
      {:ok, value} ->
        cache_entry(state.table, type, id, value)
        {:ok, value, put_known_id(state, type, id)}

      {:error, :not_found} ->
        cache_entry(state.table, type, id, @missing)
        {:ok, nil, state}

      {:error, reason} ->
        {:error, reason, state}
    end
  end

  defp cache_entry(table, type, id, value) do
    true = :ets.insert(table, {{type, id}, value})
    :ok
  end

  defp put_known_id(state, type, id) do
    update_in(state.known_ids[type], fn ids ->
      MapSet.put(ids || MapSet.new(), id)
    end)
  end

  defp drop_known_id(state, type, id) do
    update_in(state.known_ids[type], fn
      nil -> nil
      ids -> MapSet.delete(ids, id)
    end)
  end

  defp commit_with_retry(state, data), do: commit_with_retry(state, data, false)

  defp commit_with_retry(state, data, clear?) when not clear? and map_size(data) == 0,
    do: {:ok, state}

  defp commit_with_retry(state, data, clear?) do
    commit_with_retry(state, data, clear?, state.max_commit_retries)
  end

  defp commit_with_retry(state, _data, _clear?, attempts_left) when attempts_left <= 0,
    do: {:error, :commit_retry_exhausted, state}

  defp commit_with_retry(state, data, clear?, attempts_left) do
    data = normalize_data(data)
    snapshot = snapshot_for_commit(state, data, clear?)

    case apply_commit(state, data, clear?) do
      {:ok, state} ->
        {:ok, state}

      {:error, reason, state} ->
        case restore_snapshot(state, snapshot) do
          {:ok, restored_state} when attempts_left > 1 ->
            Process.sleep(restored_state.delay_between_tries_ms)
            commit_with_retry(restored_state, data, clear?, attempts_left - 1)

          {:ok, restored_state} ->
            {:error, reason, restored_state}

          {:error, rollback_reason, restored_state} ->
            {:error, {:rollback_failed, reason, rollback_reason}, restored_state}
        end
    end
  end

  defp snapshot_for_commit(state, data, false), do: snapshot_for(state, data)
  defp snapshot_for_commit(state, data, true), do: snapshot_for_clear(state, data)

  defp normalize_data(data) do
    Enum.reduce(data, %{}, fn {type, entries}, acc ->
      case normalize_entries(type, entries) do
        %{} = normalized when map_size(normalized) > 0 -> Map.put(acc, type, normalized)
        _ -> acc
      end
    end)
  end

  defp normalize_entries(_type, entries) when not is_map(entries), do: %{}

  defp normalize_entries(_type, entries) do
    Enum.reduce(entries, %{}, fn {id, value}, acc ->
      Map.put(acc, id, value)
    end)
  end

  defp apply_mutations(state, data) do
    data
    |> Enum.reduce_while(state, fn {type, entries}, acc_state ->
      apply_mutation_batch(acc_state, type, entries)
    end)
    |> case do
      {:error, reason, next_state} -> {:error, reason, next_state}
      next_state -> {:ok, next_state}
    end
  end

  defp apply_commit(state, data, false), do: apply_mutations(state, data)

  defp apply_commit(state, data, true) do
    case clear_persisted_entries(state) do
      {:ok, cleared_state} -> apply_mutations(cleared_state, data)
      {:error, reason, cleared_state} -> {:error, reason, cleared_state}
    end
  end

  defp apply_mutation_batch(state, type, entries) do
    case prepare_entries(state, type, entries) do
      {:ok, prepared_entries, next_state} ->
        case apply_prepared_entries(next_state, type, prepared_entries) do
          {:ok, updated_state} -> {:cont, updated_state}
          {:error, reason, updated_state} -> {:halt, {:error, reason, updated_state}}
        end

      {:error, reason, next_state} ->
        {:halt, {:error, reason, next_state}}
    end
  end

  defp prepare_entries(state, :"pre-key", entries), do: validate_prekey_entries(state, entries)
  defp prepare_entries(state, _type, entries), do: {:ok, entries, state}

  defp validate_prekey_entries(state, entries) do
    {deletions, updates} = Enum.split_with(entries, fn {_id, value} -> is_nil(value) end)
    updates = Map.new(updates)

    case fetch_missing(state, :"pre-key", Enum.map(deletions, &elem(&1, 0))) do
      {:ok, existing, next_state} ->
        {:ok, merge_prekey_deletions(updates, deletions, existing), next_state}

      {:error, reason, next_state} ->
        {:error, reason, next_state}
    end
  end

  defp merge_prekey_deletions(updates, deletions, existing) do
    Enum.reduce(deletions, updates, fn {id, _value}, acc ->
      if Map.has_key?(existing, id), do: Map.put(acc, id, nil), else: acc
    end)
  end

  defp apply_prepared_entries(state, type, entries) do
    Enum.reduce_while(entries, {:ok, state}, fn
      {id, nil}, {:ok, acc_state} ->
        case persistence_delete(acc_state, type, id) do
          :ok ->
            cache_entry(acc_state.table, type, id, @missing)
            {:cont, {:ok, drop_known_id(acc_state, type, id)}}

          {:error, reason} ->
            {:halt, {:error, reason, acc_state}}
        end

      {id, value}, {:ok, acc_state} ->
        case persistence_save(acc_state, type, id, value) do
          :ok ->
            cache_entry(acc_state.table, type, id, value)
            {:cont, {:ok, put_known_id(acc_state, type, id)}}

          {:error, reason} ->
            {:halt, {:error, reason, acc_state}}
        end
    end)
  end

  defp snapshot_for(state, data) do
    Enum.reduce(data, %{}, fn {type, entries}, acc ->
      Map.put(acc, type, snapshot_values_for_type(state, type, Map.keys(entries)))
    end)
  end

  defp snapshot_for_clear(state, data) do
    known_ids =
      Enum.reduce(state.known_ids, %{}, fn {type, ids}, acc ->
        Map.put(acc, type, MapSet.to_list(ids))
      end)

    ids_by_type =
      Enum.reduce(data, known_ids, fn {type, entries}, acc ->
        Map.update(acc, type, Map.keys(entries), fn existing_ids ->
          existing_ids
          |> Enum.concat(Map.keys(entries))
          |> Enum.uniq()
        end)
      end)

    Enum.reduce(ids_by_type, %{}, fn {type, ids}, acc ->
      Map.put(acc, type, snapshot_values_for_type(state, type, ids))
    end)
  end

  defp snapshot_values_for_type(state, type, ids) do
    {cached, missing_ids} = read_cached_entries(state.table, type, ids)
    fetched = fetch_snapshot_missing(state, type, missing_ids)

    Enum.reduce(ids, %{}, fn id, acc ->
      Map.put(acc, id, snapshot_value(cached, fetched, id))
    end)
  end

  defp fetch_snapshot_missing(state, type, ids) do
    case fetch_missing(state, type, ids) do
      {:ok, values, _state} -> values
      {:error, _reason, _state} -> %{}
    end
  end

  defp snapshot_value(cached, fetched, id) do
    case Map.fetch(cached, id) do
      {:ok, value} -> value
      :error -> Map.get(fetched, id, @missing)
    end
  end

  defp restore_snapshot(state, snapshot) do
    Enum.reduce_while(snapshot, {:ok, state}, fn {type, entries}, {:ok, acc_state} ->
      restore_type_entries(acc_state, type, entries)
      |> case do
        {:ok, next_state} -> {:cont, {:ok, next_state}}
        {:error, reason, next_state} -> {:halt, {:error, reason, next_state}}
      end
    end)
  end

  defp restore_type_entries(state, type, entries) do
    Enum.reduce_while(entries, {:ok, state}, fn
      {id, @missing}, {:ok, acc_state} ->
        restore_snapshot_entry(acc_state, type, id, @missing)

      {id, value}, {:ok, acc_state} ->
        restore_snapshot_entry(acc_state, type, id, value)
    end)
  end

  defp restore_snapshot_entry(state, type, id, @missing) do
    case persistence_delete(state, type, id) do
      :ok ->
        cache_entry(state.table, type, id, @missing)
        {:cont, {:ok, drop_known_id(state, type, id)}}

      {:error, reason} ->
        {:halt, {:error, reason, state}}
    end
  end

  defp restore_snapshot_entry(state, type, id, value) do
    case persistence_save(state, type, id, value) do
      :ok ->
        cache_entry(state.table, type, id, value)
        {:cont, {:ok, put_known_id(state, type, id)}}

      {:error, reason} ->
        {:halt, {:error, reason, state}}
    end
  end

  defp clear_persisted_entries(state) do
    state.known_ids
    |> Enum.reduce_while({:ok, state}, fn {type, ids}, {:ok, acc_state} ->
      clear_known_ids(acc_state, type, ids)
      |> case do
        {:ok, next_state} -> {:cont, {:ok, next_state}}
        {:error, reason, next_state} -> {:halt, {:error, reason, next_state}}
      end
    end)
    |> case do
      {:ok, next_state} ->
        true = :ets.delete_all_objects(next_state.table)
        {:ok, %{next_state | known_ids: %{}}}

      {:error, reason, next_state} ->
        {:error, reason, next_state}
    end
  end

  defp clear_known_ids(state, type, ids) do
    Enum.reduce_while(ids, {:ok, state}, fn id, {:ok, acc_state} ->
      case persistence_delete(acc_state, type, id) do
        :ok -> {:cont, {:ok, acc_state}}
        {:error, reason} -> {:halt, {:error, reason, acc_state}}
      end
    end)
  end

  defp persistence_load(state, type, id) do
    apply_persistence(state, :load_keys, [type, id], [state.persistence_context, type, id])
  end

  defp persistence_save(state, type, id, value) do
    apply_persistence(
      state,
      :save_keys,
      [type, id, value],
      [state.persistence_context, type, id, value]
    )
  end

  defp persistence_delete(state, type, id) do
    apply_persistence(state, :delete_keys, [type, id], [state.persistence_context, type, id])
  end

  defp apply_persistence(
         %{persistence_module: module, persistence_context: context},
         fun,
         args,
         ctx_args
       ) do
    _ = Code.ensure_loaded(module)

    cond do
      not is_nil(context) and function_exported?(module, fun, length(ctx_args)) ->
        apply(module, fun, ctx_args)

      function_exported?(module, fun, length(args)) ->
        apply(module, fun, args)

      true ->
        {:error, {:unsupported_persistence_operation, module, fun}}
    end
  end

  defp merge_fetched_missing(entries, _ref, _type, []), do: entries

  defp merge_fetched_missing(entries, ref, type, missing_ids) do
    case GenServer.call(ref.pid, {:fetch_missing, type, missing_ids}, :infinity) do
      {:ok, fetched} -> Map.merge(entries, fetched)
      {:error, reason} -> raise OperationError, action: :get, reason: reason
    end
  end
end