ragex

defmodule Ragex.Graph.Store do
  @moduledoc """
  Knowledge graph storage dispatcher.

  Manages nodes (modules, functions, types, etc.) and edges (calls, imports, etc.)
  representing relationships in the codebase.

  The actual storage is delegated to the configured backend module
  (see `Ragex.Store.Backend`). The default is ETS; dllb can be selected via:

      config :ragex, :store_backend, :dllb
  """

  use GenServer
  require Logger

  alias Ragex.Embeddings.{FileTracker, Persistence}
  alias Ragex.Graph.Persistence, as: GraphPersistence
  alias Ragex.Store.Backend

  @nodes_table :ragex_nodes
  @edges_table :ragex_edges
  @embeddings_table :ragex_embeddings

  _timeout =
    :ragex
    |> Application.compile_env(:timeouts, [])
    |> Keyword.get(:store, :infinity)

  @functions_limit Application.compile_env(:ragex, :functions_limit, 1_000)
  @nodes_limit Application.compile_env(:ragex, :nodes_limit, 1_000)
  @embeddings_limit Application.compile_env(:ragex, :embeddings_limit, 1_000)

  # Client API

  def start_link(opts \\ []) do
    GenServer.start_link(__MODULE__, opts, name: __MODULE__)
  end

  @doc """
  Adds a node to the graph.

  Node types: :module, :function, :type, :variable, :file
  """
  def add_node(node_type, node_id, data) do
    GenServer.cast(__MODULE__, {:add_node, node_type, node_id, data})
  end

  @doc """
  Retrieves a node by its composite key.

  ## Parameters
  - `node_key` - Tuple of `{node_type, node_id}`

  ## Returns
  - Node data map if found
  - `nil` if not found

  ## Examples

      iex> Store.add_node(:module, MyModule, %{name: MyModule})
      iex> Store.get_node({:module, MyModule})
      %{name: MyModule}

      iex> Store.get_node({:module, NonExistent})
      nil
  """
  def get_node({node_type, node_id}) do
    find_node(node_type, node_id)
  end

  @doc """
  Finds a node by type and id.
  """
  def find_node(node_type, node_id) do
    backend().find_node(node_type, node_id)
  end

  @doc """
  Finds a function node by module and name (any arity).
  """
  def find_function(module, name) do
    backend().find_function(module, name)
  end

  @doc """
  Retrieves a function node by module, name, and arity.

  Returns the function data or nil if not found.

  ## Parameters
  - `module` - The module containing the function
  - `name` - The function name (atom)
  - `arity` - The function arity (non-negative integer)

  ## Examples

      iex> Store.add_node(:function, {MyModule, :test, 2}, %{name: :test, arity: 2})
      iex> Store.get_function(MyModule, :test, 2)
      %{name: :test, arity: 2}

      iex> Store.get_function(MyModule, :nonexistent, 1)
      nil
  """
  def get_function(module, name, arity) do
    find_node(:function, {module, name, arity})
  end

  @doc """
  Retrieves a module node by name.

  Returns the module data or nil if not found.

  ## Examples

      iex> Store.add_node(:module, MyModule, %{name: MyModule, file: "lib/my_module.ex"})
      iex> Store.get_module(MyModule)
      %{name: MyModule, file: "lib/my_module.ex"}

      iex> Store.get_module(NonExistentModule)
      nil
  """
  def get_module(module) do
    find_node(:module, module)
  end

  @doc """
  Lists all module nodes.

  ## Returns
  List of module maps with keys:
  - `:id` - Module identifier (atom)
  - `:data` - Module data

  ## Examples

      iex> Store.add_node(:module, ModuleA, %{name: ModuleA, file: "lib/a.ex"})
      iex> Store.add_node(:module, ModuleB, %{name: ModuleB, file: "lib/b.ex"})
      iex> Store.list_modules()
      [%{id: ModuleA, data: %{name: ModuleA, file: "lib/a.ex"}}, ...]
  """
  def list_modules do
    :module
    |> backend().list_nodes(:infinity)
    |> Enum.map(&%{id: &1.id, data: &1.data})
  end

  @doc """
  Lists function nodes with optional filtering by module.

  ## Parameters
  - `opts` - Keyword list with options:
    - `:module` - Filter by module (optional)
    - `:limit` - Maximum number of functions to return (default: 1000)

  ## Returns
  List of function maps with keys:
  - `:id` - Function identifier tuple `{module, name, arity}`
  - `:data` - Function data

  ## Examples

      iex> Store.add_node(:function, {MyModule, :test, 2}, %{name: :test, arity: 2})
      iex> Store.list_functions()
      [%{id: {MyModule, :test, 2}, data: %{name: :test, arity: 2}}]

      iex> Store.list_functions(module: MyModule, limit: 50)
      [%{id: {MyModule, :test, 2}, data: ...}, ...]
  """
  def list_functions(opts \\ []) do
    module = Keyword.get(opts, :module)
    limit = Keyword.get(opts, :limit, @functions_limit)

    nodes =
      case module do
        nil ->
          backend().list_nodes(:function, limit)

        mod ->
          :function
          |> backend().list_nodes(:infinity)
          |> Enum.filter(fn
            %{id: {m, _, _}} -> m == mod
            _ -> false
          end)
          |> Enum.take(limit)
      end

    Enum.map(nodes, &%{id: &1.id, data: &1.data})
  end

  @doc """
  Lists nodes with optional filtering by type.
  """
  def list_nodes(node_type \\ nil, limit \\ @nodes_limit) do
    backend().list_nodes(node_type, limit)
  end

  @doc """
  Counts nodes of a specific type.
  """
  def count_nodes_by_type(node_type) do
    backend().count_nodes_by_type(node_type)
  end

  @doc """
  Adds an edge between two nodes.

  Edge types: :calls, :imports, :defines, :inherits, :implements

  ## Options
  - `:weight` - Edge weight (default: 1.0) for weighted graph algorithms
  - `:metadata` - Additional metadata map
  """
  def add_edge(from_node, to_node, edge_type, opts \\ []) do
    GenServer.cast(__MODULE__, {:add_edge, from_node, to_node, edge_type, opts})
  end

  @doc """
  Gets all outgoing edges from a node of a specific type.
  """
  def get_outgoing_edges(from_node, edge_type) do
    backend().get_outgoing_edges(from_node, edge_type)
  end

  @doc """
  Gets all incoming edges to a node of a specific type.
  """
  def get_incoming_edges(to_node, edge_type) do
    backend().get_incoming_edges(to_node, edge_type)
  end

  @doc """
  Gets the weight of an edge between two nodes.

  Returns the weight (default: 1.0) if edge exists, nil otherwise.
  """
  def get_edge_weight(from_node, to_node, edge_type) do
    backend().get_edge_weight(from_node, to_node, edge_type)
  end

  @doc """
  Lists all edges with optional filtering by type and limit.

  ## Parameters
  - `opts` - Keyword list with options:
    - `:edge_type` - Filter by edge type (optional)
    - `:limit` - Maximum number of edges to return (default: 1000)

  ## Returns
  List of edge maps with keys:
  - `:from` - Source node
  - `:to` - Target node
  - `:type` - Edge type
  - `:metadata` - Edge metadata including weight

  ## Examples

      iex> Store.list_edges(limit: 100)
      [%{from: node1, to: node2, type: :calls, metadata: %{weight: 1.0}}, ...]

      iex> Store.list_edges(edge_type: :imports)
      [%{from: mod1, to: mod2, type: :imports, metadata: %{weight: 1.0}}, ...]
  """
  def list_edges(opts \\ []) do
    backend().list_edges(opts)
  end

  @doc """
  Removes a node from the graph.

  Also removes all edges connected to this node (both incoming and outgoing)
  and any embedding associated with the node.

  Returns `:ok` if successful, `{:error, :timeout}` on timeout.
  """
  def remove_node(node_type, node_id) do
    GenServer.cast(__MODULE__, {:remove_node, node_type, node_id})
  end

  @doc """
  Merges additional metadata into an existing node.

  If the node exists, its metadata map is deep-merged with `new_metadata`.
  If the node does not exist, this is a no-op.

  ## Parameters

  - `node_type` -- `:module`, `:function`, etc.
  - `node_id` -- the node identifier
  - `new_metadata` -- map of metadata to merge
  """
  def update_node_metadata(node_type, node_id, new_metadata) when is_map(new_metadata) do
    backend().update_node_metadata(node_type, node_id, new_metadata)
  end

  @doc """
  Clears all data from the graph.
  """
  def clear do
    GenServer.cast(__MODULE__, :clear)
  end

  @doc """
  Switches the store to a specific project path.

  Clears all current graph data (nodes, edges, embeddings, file tracker),
  then loads the cached graph and embeddings for the given project path.
  This ensures the store contains only data for the target project.

  ## Parameters

  - `project_path` - Absolute path to the project directory

  ## Returns

  - `:ok`
  """
  @spec load_project(String.t()) :: :ok
  def load_project(project_path) do
    GenServer.call(__MODULE__, {:load_project, project_path}, :infinity)
  end

  @doc """
  Stores an embedding vector for a node.
  """
  def store_embedding(node_type, node_id, embedding, text) do
    GenServer.cast(__MODULE__, {:store_embedding, node_type, node_id, embedding, text})
  end

  @doc """
  Retrieves the embedding for a node.

  Returns `{embedding, text}` tuple or `nil` if not found.
  """
  def get_embedding(node_type, node_id) do
    backend().get_embedding(node_type, node_id)
  end

  @doc """
  Lists all embeddings with optional type filter.

  Returns list of `{node_type, node_id, embedding, text}` tuples.
  """
  def list_embeddings(node_type \\ nil, limit \\ @embeddings_limit) do
    backend().list_embeddings(node_type, limit)
  end

  @doc """
  Synchronous barrier: blocks until all pending casts have been processed.

  Useful in tests to ensure async operations (add_node, add_edge, etc.)
  have been applied to ETS before querying.
  """
  def sync do
    GenServer.call(__MODULE__, :sync)
  end

  @doc """
  Returns statistics about the graph.
  """
  def stats do
    backend().stats()
  end

  @doc """
  Returns the ETS table reference for embeddings.

  Useful for direct access or persistence operations.
  """
  def embeddings_table, do: @embeddings_table

  # Server Callbacks

  @impl true
  def init(_opts) do
    # Create ETS tables
    :ets.new(@nodes_table, [:named_table, :set, :public, read_concurrency: true])
    :ets.new(@edges_table, [:named_table, :bag, :public, read_concurrency: true])
    :ets.new(@embeddings_table, [:named_table, :set, :public, read_concurrency: true])

    # Initialize file tracker for incremental updates
    FileTracker.init()

    # At startup, load CWD-based cache (backward compat for MCP server / interactive use).
    # Agent.Core.analyze_project will call load_project/1 to switch to the correct path.
    do_load_project_cache(nil)

    {:ok, %{project_path: nil}}
  end

  @impl true
  def handle_call(:sync, _from, state) do
    {:reply, :ok, state}
  end

  @impl true
  def handle_call({:load_project, project_path}, _from, state) do
    # Clear all tables
    :ets.delete_all_objects(@nodes_table)
    :ets.delete_all_objects(@edges_table)
    :ets.delete_all_objects(@embeddings_table)
    FileTracker.clear_all()

    # Load caches for the target project path
    do_load_project_cache(project_path)

    {:reply, :ok, %{state | project_path: project_path}}
  end

  @impl true
  def handle_cast({:add_node, node_type, node_id, data}, state) do
    backend().store_node(node_type, node_id, data)
    {:noreply, state}
  end

  @impl true
  def handle_cast({:add_edge, from_node, to_node, edge_type, opts}, state) do
    backend().store_edge(from_node, to_node, edge_type, opts)
    {:noreply, state}
  end

  @impl true
  def handle_cast({:store_embedding, node_type, node_id, embedding, text}, state) do
    backend().store_embedding(node_type, node_id, embedding, text)
    {:noreply, state}
  end

  @impl true
  def handle_cast({:remove_node, node_type, node_id}, state) do
    backend().remove_node(node_type, node_id)
    {:noreply, state}
  end

  @impl true
  def handle_cast(:clear, state) do
    backend().clear()
    {:noreply, state}
  end

  @impl true
  def terminate(reason, state) do
    # Save to disk on normal shutdown, using the stored project path
    if reason == :shutdown or reason == :normal do
      pp = state[:project_path]

      case Persistence.save(@embeddings_table, pp) do
        {:ok, path} ->
          Logger.info("Embeddings saved to #{path}")

        {:error, err} ->
          Logger.error("Failed to save embeddings: #{inspect(err)}")
      end

      case GraphPersistence.save(pp) do
        {:ok, path} ->
          Logger.info("Graph saved to #{path}")

        {:error, err} ->
          Logger.error("Failed to save graph: #{inspect(err)}")
      end
    else
      Logger.warning("Graph store terminating abnormally: #{inspect(reason)}, skipping save")
    end

    # ETS tables are automatically cleaned up
    :ok
  end

  # Loads graph and embedding caches for a given project path (nil = CWD).
  defp do_load_project_cache(project_path) do
    # For ETS backend, load from disk cache
    if Backend.module() == Ragex.Store.Backend.ETS do
      case Persistence.load(project_path) do
        {:ok, count} ->
          Logger.info("Loaded #{count} cached embeddings")

        {:error, :not_found} ->
          Logger.debug("No embedding cache found")

        {:error, :incompatible} ->
          Logger.warning("Embedding cache incompatible with current model")

        {:error, reason} ->
          Logger.warning("Failed to load embedding cache: #{inspect(reason)}")
      end

      case GraphPersistence.load(project_path) do
        {:ok, %{nodes: n, edges: e}} ->
          Logger.info("Loaded graph from cache: #{n} nodes, #{e} edges")

        {:error, :not_found} ->
          Logger.debug("No graph cache found")

        {:error, reason} ->
          Logger.warning("Failed to load graph cache: #{inspect(reason)}")
      end
    end

    # Bootstrap backend (no-op for ETS, schema setup for dllb)
    case backend().bootstrap() do
      :ok -> Logger.info("Store backend bootstrapped (#{Backend.module()})")
      {:error, reason} -> Logger.warning("Store backend bootstrap failed: #{inspect(reason)}")
    end
  end

  defp backend, do: Backend.module()
end