localize_mf2_treesitter

defmodule Localize.Mf2.TreeSitter.Query do
  @moduledoc """
  Tree-sitter query compilation and execution.

  A query is a set of [S-expression patterns](https://tree-sitter.github.io/tree-sitter/using-parsers#pattern-matching-with-queries)
  compiled against the MF2 grammar. Running a query against a node
  yields either matches (with pattern provenance) or a flat list of
  captures.

  Two queries ship with this package, mirroring the ones in the
  grammar repo:

  * `:highlights` — the editor-highlight captures.
  * `:injections` — host-language injection hints (for `~M` sigil etc).

  Load them with `load/1` and run against any `Node` you have.
  """

  alias Localize.Mf2.TreeSitter.{Nif, Node}

  @enforce_keys [:ref]
  defstruct [:ref]

  @type t :: %__MODULE__{ref: reference()}

  @type capture :: {name :: binary(), Node.t()}
  @type match :: %{pattern_index: non_neg_integer(), captures: [capture()]}

  @bundled %{
    highlights: "highlights.scm",
    injections: "injections.scm"
  }

  @doc """
  Compile a query from its S-expression source.

  ### Arguments

  * `source` is a binary containing one or more tree-sitter query
    patterns.

  ### Returns

  * `{:ok, query}` on success.

  * `{:error, {:query_error, byte_offset, kind}}` if the query is
    malformed. `kind` is one of `:syntax`, `:node_type`, `:field`,
    `:capture`, `:structure`, `:language`, or `:unknown`, matching
    tree-sitter's `TSQueryError` enum.

  ### Examples

      iex> {:ok, query} = Localize.Mf2.TreeSitter.Query.new("(variable (name) @var)")
      iex> Localize.Mf2.TreeSitter.Query.capture_names(query)
      ["var"]

  """
  @spec new(binary()) :: {:ok, t()} | {:error, {:query_error, non_neg_integer(), atom()}}
  def new(source) when is_binary(source) do
    case Nif.query_new(source) do
      {:ok, ref} -> {:ok, %__MODULE__{ref: ref}}
      {:error, reason} -> {:error, reason}
    end
  end

  @doc """
  Load one of the bundled queries shipped under `priv/queries/`.

  ### Arguments

  * `name` is `:highlights`, `:injections`, or a filename string
    resolved relative to `priv/queries/`.

  ### Returns

  * `{:ok, query}` on success.
  * `{:error, {:file, posix_reason}}` if the file cannot be read.
  * `{:error, {:query_error, offset, kind}}` if the query does not compile.

  ### Examples

      iex> {:ok, _query} = Localize.Mf2.TreeSitter.Query.load(:highlights)

  """
  @spec load(atom() | String.t()) :: {:ok, t()} | {:error, term()}
  def load(name) when is_atom(name) or is_binary(name) do
    filename =
      case name do
        n when is_atom(n) ->
          Map.get(@bundled, n) || "#{n}.scm"

        n when is_binary(n) ->
          n
      end

    path =
      :localize_mf2_treesitter
      |> :code.priv_dir()
      |> Path.join(Path.join("queries", filename))

    case File.read(path) do
      {:ok, source} -> new(source)
      {:error, reason} -> {:error, {:file, reason}}
    end
  end

  @doc "Return the number of patterns compiled into the query."
  @spec pattern_count(t()) :: non_neg_integer()
  def pattern_count(%__MODULE__{ref: ref}), do: Nif.query_pattern_count(ref)

  @doc "Return the number of distinct capture names in the query."
  @spec capture_count(t()) :: non_neg_integer()
  def capture_count(%__MODULE__{ref: ref}), do: Nif.query_capture_count(ref)

  @doc "Return the list of capture names in the order they first appear in the query."
  @spec capture_names(t()) :: [binary()]
  def capture_names(%__MODULE__{ref: ref}), do: Nif.query_capture_names(ref)

  @doc """
  Run the query against a node and return a list of matches.

  Each match is `%{pattern_index: i, captures: [{name, node}]}`. Use
  `captures/2` if pattern provenance is not needed.

  ### Arguments

  * `query` is a `#{inspect(__MODULE__)}`.
  * `node` is a `Localize.Mf2.TreeSitter.Node`.

  ### Returns

  * A list of match maps, in match order.

  ### Examples

      iex> {:ok, tree} = Localize.Mf2.TreeSitter.parse("hello {$name}")
      iex> root = Localize.Mf2.TreeSitter.root(tree)
      iex> {:ok, query} = Localize.Mf2.TreeSitter.Query.new("(variable (name) @var)")
      iex> [match] = Localize.Mf2.TreeSitter.Query.matches(query, root)
      iex> [{"var", node}] = match.captures
      iex> Localize.Mf2.TreeSitter.Node.text(node)
      "name"

  """
  @spec matches(t(), Node.t()) :: [match()]
  def matches(%__MODULE__{ref: qref}, %Node{handle: h, source: src}) do
    qref
    |> Nif.query_matches(h)
    |> Enum.map(fn %{pattern_index: idx, captures: caps} ->
      %{pattern_index: idx, captures: wrap_captures(caps, src)}
    end)
  end

  @doc """
  Run the query and return a flat list of `{name, node}` captures.

  This mirrors tree-sitter's `next_capture` iteration — useful for
  editor highlight passes where pattern provenance is irrelevant.

  ### Arguments

  * `query` is a `#{inspect(__MODULE__)}`.
  * `node` is a `Localize.Mf2.TreeSitter.Node`.

  ### Returns

  * A list of `{name, Node.t()}` tuples, in capture order.

  ### Examples

      iex> {:ok, tree} = Localize.Mf2.TreeSitter.parse("hello {$name}")
      iex> root = Localize.Mf2.TreeSitter.root(tree)
      iex> {:ok, query} = Localize.Mf2.TreeSitter.Query.load(:highlights)
      iex> caps = Localize.Mf2.TreeSitter.Query.captures(query, root)
      iex> Enum.any?(caps, fn {name, _} -> name == "variable" end)
      true

  """
  @spec captures(t(), Node.t()) :: [capture()]
  def captures(%__MODULE__{ref: qref}, %Node{handle: h, source: src}) do
    qref
    |> Nif.query_captures(h)
    |> wrap_captures(src)
  end

  defp wrap_captures(raw_captures, src) do
    Enum.map(raw_captures, fn {name, handle} ->
      {name, %Node{handle: handle, source: src}}
    end)
  end
end