jsonpath

defmodule JsonPath do
  @moduledoc """
  Minimal JSONPath implementation for Elixir.

  This module provides a simple JSONPath engine with:

    * Tokenization and parsing using LEEX and YECC generated Erlang modules.
    * AST evaluation against Elixir maps, lists, and primitives.
    * Full support for child, descendant, wildcard, slices, unions, and basic filters.
    * Convenience helper functions to query a map/list without manually parsing or tokenizing.

  ## Example

      iex> data = %{"store" => %{"book" => [%{"price" => 10}, %{"price" => 20}]}}
      iex> JsonPath.query(data, "$.store.book[*].price")
      [{"$['store']['book'][0]['price']", 10}, {"$['store']['book'][1]['price']", 20}]

  """

  @typedoc "JSONPath AST generated by the parser"
  @type ast :: {:jsonpath, :root, list(any())} | map() | {:path, list(any())}

  @typedoc "A node path in the JSON structure"
  @type path :: list(String.t() | integer())

  @typedoc "Result of evaluating a JSONPath AST: list of {path_string, value}"
  @type eval_result :: list({String.t(), any()})

  @typedoc "Result of tokenize/1"
  @type tokenize_result :: {:ok, list(any()), integer()} | {:error, any()}

  @typedoc "Result of parse/1"
  @type parse_result :: {:ok, ast()} | {:error, any()}

  @doc """
  Tokenizes a JSONPath query string into a list of tokens using the LEEX lexer.

  Returns `{:ok, tokens, line}` or `{:error, reason}`.

  ## Parameters

    * `query` - JSONPath string, e.g. "$.store.book[*].price"

  """
  @spec tokenize(String.t()) :: tokenize_result()
  def tokenize(query) when is_binary(query) do
    trimmed = String.trim(query)

    if trimmed != query do
      {:error, "JSONPath expressions cannot have leading or trailing whitespace"}
    else
      ensure_loaded(:jsonpath_lexer)
      :jsonpath_lexer.string(String.to_charlist(query))
    end
  end

  @doc """
  Parses a list of tokens into an AST using the YECC parser.

  Returns `{:ok, ast}` or `{:error, reason}`.
  """
  @spec parse(list(any())) :: parse_result()
  def parse(tokens) when is_list(tokens) do
    try do
      ensure_loaded(:jsonpath_parser)
      :jsonpath_parser.parse(tokens)
    rescue
      e in CaseClauseError -> {:error, e.term}
      e -> {:error, e.term}
    end
  end

  @doc """
  Evaluate a JSONPath AST against data.

  ## Parameters

    * `ast` - the parsed JSONPath AST
    * `data` - Elixir map or list to evaluate against

  ## Returns

    * list of tuples `{path_string, value}`
  """
  @spec evaluate(ast(), map() | list()) :: eval_result() | {:error, any()}
  def evaluate(ast, data) do
    case ast do
      {:jsonpath, :root, segments} when is_list(segments) -> traverse([{["$"], data}], segments)
      %{root: :root, segments: segs} when is_list(segs) -> traverse([{["$"], data}], segs)
      %{segments: segs} when is_list(segs) -> traverse([{["$"], data}], segs)
      {:path, segs} when is_list(segs) -> traverse([{["$"], data}], segs)
      other -> {:error, {:invalid_ast, other}}
    end
  end

  @doc """
  Convenience function: parse and evaluate a JSONPath query string in one call.

  ## Example

      iex> JsonPath.query(%{"a" => 1}, "$.a")
      [{"$['a']", 1}]
  """
  @spec query(map() | list(), String.t()) :: eval_result() | {:error, any()}
  def query(data, path_string) when is_binary(path_string) and (is_map(data) or is_list(data)) do
    with {:ok, tokens, _line} <- tokenize(path_string),
         {:ok, ast} <- parse(tokens) do
      evaluate(ast, data)
    else
      {:error, reason} -> {:error, reason}
    end
  end

  ## PRIVATE HELPERS

  @spec ensure_loaded(module()) :: :ok
  defp ensure_loaded(module) do
    case :code.is_loaded(module) do
      {:file, _} -> :ok
      _ -> :code.load_file(module)
    end
  end

  @spec traverse(list({path(), any()}), list(any())) :: eval_result()
  defp traverse(nodes, []), do: Enum.map(nodes, fn {p, v} -> {path_join(p), v} end)

  defp traverse(nodes, [seg | rest]) do
    expanded =
      Enum.flat_map(nodes, fn {path, node} ->
        apply_segment(seg, {path, node})
      end)

    traverse(expanded, rest)
  end

  # apply child or descendant segment
  defp apply_segment({child_type, selectors}, {path, node}) do
    case child_type do
      :child ->
        Enum.flat_map(selectors, fn sel -> apply_selector(sel, {path, node}, :child) end)

      :descendant ->
        # descendant: node itself then recursively into children
        nodes = collect_descendants({path, node})

        Enum.flat_map(nodes, fn n ->
          Enum.flat_map(selectors, &apply_selector(&1, n, :descendant))
        end)
    end
  end

  defp collect_descendants({path, value}) do
    # include the node itself
    base = [{path, value}]

    children =
      case value do
        %{} = m ->
          Enum.flat_map(Map.to_list(m), fn {k, v} ->
            collect_descendants({path ++ [to_string(k)], v})
          end)

        l when is_list(l) ->
          Enum.with_index(l)
          |> Enum.flat_map(fn {v, idx} ->
            collect_descendants({path ++ [idx], v})
          end)

        _ ->
          []
      end

    base ++ children
  end

  defp apply_selector({:name, name}, {path, node}, _mode) when is_binary(name) do
    case node do
      %{} = m ->
        case Map.fetch(m, name) do
          {:ok, v} -> [{path ++ [name], v}]
          :error -> []
        end

      _ ->
        []
    end
  end

  defp apply_selector({:name, name}, {path, node}, _mode) when is_atom(name) do
    apply_selector({:name, to_string(name)}, {path, node}, :child)
  end

  defp apply_selector({:wildcard}, {path, node}, _mode) do
    case node do
      %{} = m ->
        Enum.map(Map.to_list(m), fn {k, v} -> {path ++ [to_string(k)], v} end)

      l when is_list(l) ->
        Enum.with_index(l) |> Enum.map(fn {v, i} -> {path ++ [i], v} end)

      _ ->
        []
    end
  end

  defp apply_selector({:index, idx}, {path, node}, _mode) when is_integer(idx) do
    case node do
      l when is_list(l) ->
        n = length(l)
        real_i = if idx < 0, do: n + idx, else: idx
        if real_i >= 0 and real_i < n, do: [{path ++ [real_i], Enum.at(l, real_i)}], else: []

      _ ->
        []
    end
  end

  defp apply_selector({:slice, slice}, {path, node}, _mode) do
    case node do
      l when is_list(l) ->
        indices = normalize_slice(length(l), slice)
        Enum.map(indices, fn i -> {path ++ [i], Enum.at(l, i)} end)

      _ ->
        []
    end
  end

  defp apply_selector({:union, items}, {path, node}, mode) do
    Enum.flat_map(items, fn item ->
      apply_selector(item, {path, node}, mode)
    end)
  end

  defp apply_selector({:filter, expr}, {path, node}, _mode) do
    # filter applies to arrays/objects; for objects, iterate its children; for arrays, iterate elements
    case node do
      l when is_list(l) ->
        Enum.with_index(l)
        |> Enum.flat_map(fn {v, i} ->
          if eval_filter(expr, v), do: [{path ++ [i], v}], else: []
        end)

      %{} = m ->
        Enum.flat_map(Map.to_list(m), fn {k, v} ->
          if eval_filter(expr, v), do: [{path ++ [to_string(k)], v}], else: []
        end)

      _ ->
        []
    end
  end

  # slice normalization: accepts slice shape produced by parser
  defp normalize_slice(len, {:start_end_step, s, e, step}) do
    range_from_slice(len, s, e, step)
  end

  defp normalize_slice(len, {:start_end, s, e}), do: range_from_slice(len, s, e, 1)
  defp normalize_slice(len, {:start_omitted_end, e}), do: range_from_slice(len, 0, e, 1)
  defp normalize_slice(len, {:start_end_omitted, s}), do: range_from_slice(len, s, len, 1)

  defp normalize_slice(len, {:start_omitted_end_step, step}),
    do: range_from_slice(len, 0, len, step)

  defp normalize_slice(_len, {:omitted_all}), do: []

  defp range_from_slice(len, start, stop, step) do
    st = if start < 0, do: max(len + start, 0), else: min(start, len)
    sp = if stop < 0, do: max(len + stop, 0), else: min(stop, len)

    cond do
      step == 0 ->
        []

      step > 0 ->
        if st >= sp do
          []
        else
          st..(sp - 1)//step |> Enum.to_list()
        end

      step < 0 ->
        if st <= sp do
          []
        else
          st..(sp + 1)//step |> Enum.to_list()
        end
    end
  end

  ## FILTER EVAL: evaluation of filter AST nodes against a value.
  ## Basic implementation: literals, existence tests via `@.name` are supported, comparisons.

  defp eval_filter({:or, a, b}, node), do: eval_filter(a, node) or eval_filter(b, node)
  defp eval_filter({:and, a, b}, node), do: eval_filter(a, node) and eval_filter(b, node)
  defp eval_filter({:not, a}, node), do: not eval_filter(a, node)

  defp eval_filter({:cmp, op, left, right}, node) do
    l = eval_primary(left, node)
    r = eval_primary(right, node)
    compare_values(op, l, r)
  end

  defp eval_filter({:query, :relative, qsegs}, node) do
    # run a tiny singular query starting at node; returns truthy if selects >= 1
    results = run_singular_query(node, qsegs)
    length(results) > 0
  end

  defp eval_filter({:query, :absolute, _}, _node) do
    # absolute ($) not implemented in filter context for demo; return false
    false
  end

  defp eval_primary({:lit, v}, _node), do: v

  defp eval_primary({:query, :relative, qsegs}, node) do
    # return first literal value if query returns one element
    case run_singular_query(node, qsegs) do
      [{_path, val} | _] -> val
      [] -> nil
    end
  end

  defp compare_values(:eq, a, b), do: a == b
  defp compare_values(:ne, a, b), do: a != b
  defp compare_values(:lt, a, b) when is_number(a) and is_number(b), do: a < b
  defp compare_values(:le, a, b) when is_number(a) and is_number(b), do: a <= b
  defp compare_values(:gt, a, b) when is_number(a) and is_number(b), do: a > b
  defp compare_values(:ge, a, b) when is_number(a) and is_number(b), do: a >= b
  defp compare_values(_, _, _), do: false

  defp run_singular_query(node, qsegs) do
    # qsegs is list of {:qname, name} or {:qindex, idx}
    # start at node with path []
    walk_query([{[], node}], qsegs)
  end

  defp walk_query(nodes, []), do: nodes

  defp walk_query(nodes, [seg | rest]) do
    next =
      Enum.flat_map(nodes, fn {path, node} ->
        case seg do
          {:qname, name} ->
            case node do
              %{} = m ->
                name_str = to_string(name)

                case Map.fetch(m, name_str) do
                  {:ok, v} -> [{path ++ [name_str], v}]
                  :error -> []
                end

              _ ->
                []
            end

          {:qindex, idx} when is_integer(idx) ->
            case node do
              l when is_list(l) ->
                n = length(l)
                i = if idx < 0, do: n + idx, else: idx
                if i >= 0 and i < n, do: [{path ++ [i], Enum.at(l, i)}], else: []

              _ ->
                []
            end
        end
      end)

    walk_query(next, rest)
  end

  defp path_join(parts) do
    "$" <>
      Enum.map_join(tl(parts), "", fn
        part when is_integer(part) -> "[#{part}]"
        part when is_binary(part) -> "['#{part}']"
      end)
  end
end