defmodule KreuzbergEmbeddingBackendBridge do
  @moduledoc """
  GenServer bridge for EmbeddingBackend implementation in kreuzberg.

  Handles incoming trait method calls from Rust and dispatches them to an implementation module.
  """

  use GenServer

  require Logger

  @doc """
  Start a GenServer linked to the current process.

  impl_module should be a module that implements the EmbeddingBackend trait methods.
  """
  def start_link(impl_module) do
    GenServer.start_link(__MODULE__, impl_module, name: __MODULE__)
  end

  @impl GenServer
  def init(impl_module) do
    {:ok, impl_module}
  end

  @doc """
  Handle an incoming trait call message.

  Message format: {:trait_call, method_atom, args_json, reply_id}
  """
  @impl GenServer
  def handle_info({:trait_call, method, args_json, reply_id}, impl_module) do
    try do
      args = Jason.decode!(args_json)
      method_name = to_string(method)
      ordered_args = ordered_args(impl_module, method_name, args)

      # Dispatch to the implementation module
      result = apply(impl_module, String.to_existing_atom(method_name), ordered_args)

      # Send result back to Rust
      Kreuzberg.Native.complete_trait_call(reply_id, Jason.encode!(result))
    rescue
      e ->
        Logger.error("Error calling {impl_module}.{method}: {Exception.message(e)}")
        Kreuzberg.Native.fail_trait_call(reply_id, Exception.message(e))
    end

    {:noreply, impl_module}
  end

  defp ordered_args(impl_module, method_name, args) when is_map(args) do
    if function_exported?(impl_module, :__alef_arg_order__, 1) do
      impl_module.__alef_arg_order__(method_name)
      |> Enum.map(&Map.fetch!(args, &1))
    else
      args
      |> Map.keys()
      |> Enum.sort()
      |> Enum.map(&Map.fetch!(args, &1))
    end
  end

  defp ordered_args(_impl_module, _method_name, args) when is_list(args), do: args

  @doc """
  Register an implementation module, starting a GenServer to handle trait calls.
  """
  def register(impl_module) do
    plugin_name = impl_module.name()
    {:ok, pid} = start_link(impl_module)
    Kreuzberg.Native.register_embedding_backend(pid, plugin_name)
  end
end
