jsonpath_ex

defmodule JSONPathEx.Evaluator do
  @moduledoc """
  Evaluates JSONPath Abstract Syntax Trees (ASTs) against JSON data.

  Supports JSONPath features such as filters, recursive descent, array slicing,
  logical operators, arithmetic, and built-in functions.
  """

  @comparison_operators [:<, :>, :<=, :>=, :==, :!=, :===, :in, :"not in"]
  @ordering_operators [:<, :>, :<=, :>=]
  @arithmetic_operators [:+, :-, :*, :/, :%]

  # Sentinel used inside filter expressions to distinguish "no value" (missing
  # key, missing index, type mismatch on traversal) from an explicit `null`.
  # Comparisons involving this value follow RFC 9535 "Nothing" semantics.
  @missing :"$$__jsonpath_ex_missing__$$"
  def missing_marker, do: @missing

  # Internal default for Enum.at/3 when we need to detect "out of range".
  @no_default :"$$__jsonpath_ex_no_default__$$"

  @doc """
  Evaluates a JSONPath AST against the provided JSON data.

  Returns the evaluation result or an empty list if the path doesn't match.
  """
  def evaluate(ast, json) do
    evaluate1(ast, json, json) |> demarker()
  end

  defp evaluate1(ast, json, original_json) do
    ast = preprocess(ast)
    Enum.reduce(ast, json, fn node, acc -> eval_ast(node, acc, original_json) end)
  end

  # Convert the missing sentinel to nil at user-visible boundaries.
  defp demarker(@missing), do: nil
  defp demarker(value) when is_list(value), do: Enum.map(value, &demarker/1)
  defp demarker(value) when is_map(value), do: value
  defp demarker(value), do: value

  # ---------------------------------------------------------------------------
  # Preprocessing: merge standalone {:deep_scan, ".."} with the node that
  # follows it so the evaluator can handle them as a single operation.
  # ---------------------------------------------------------------------------
  defp preprocess([{:deep_scan, ".."}, {:bracket_child, keys} | rest]) do
    [{:deep_scan_bracket, keys} | preprocess(rest)]
  end

  defp preprocess([{:deep_scan, ".."}, {:array, array_op} | rest]) do
    [{:deep_scan_array, array_op} | preprocess(rest)]
  end

  # $...key  →  same AST shape the parser produces for $..key;
  # re-feed through preprocess so deep_scan_key can grab the next node.
  defp preprocess([{:deep_scan, ".."}, {:dot_child, [key]} | rest]) when is_binary(key) do
    preprocess([{:dot_child, [{:deep_scan, ".."}, key]} | rest])
  end

  # Trailing $.. with nothing after it
  defp preprocess([{:deep_scan, ".."} | rest]) do
    [{:deep_scan_standalone, ".."} | preprocess(rest)]
  end

  # $..key followed by another step (array, filter, function, child …)
  # → apply that step to each scan result individually.
  defp preprocess([{:dot_child, [{:deep_scan, _}, key]}, next | rest]) do
    [{:deep_scan_key, {key, next}} | preprocess(rest)]
  end

  defp preprocess([node | rest]), do: [node | preprocess(rest)]
  defp preprocess([]), do: []

  # ---------------------------------------------------------------------------
  # Core AST dispatch
  # ---------------------------------------------------------------------------
  defp eval_ast(ast, json, original_json) do
    case ast do
      {:root, _} ->
        original_json

      {:dot, _} ->
        json

      # .* on a map → its values; on a list → identity
      {:dot_child, [dot_wildcard: _]} ->
        if is_map(json), do: Map.values(json), else: json

      # ..* → all descendants
      {:dot_child, [deep_scan_wildcard: _]} ->
        deepscan(json)

      # ..key (produced by parser directly, or by preprocess for $...key)
      {:dot_child, [{:deep_scan, _}, key]} ->
        scan(json, key)

      # Single-key dot or bracket child
      {child_key, [key]} when child_key in [:dot_child, :bracket_child] ->
        get(json, key)

      # Multi-key bracket child  $['a','b']
      {:bracket_child, keys} ->
        multi_key_bracket(json, keys)

      {:wildcard, _} ->
        json

      {:recursive_descent, _} ->
        json

      # Merged deep_scan nodes (from preprocess)
      {:deep_scan_key, {key, next_op}} ->
        scan(json, key)
        |> Enum.flat_map(fn node ->
          result = eval_ast(next_op, node, original_json)
          to_list_value(result)
        end)

      {:deep_scan_bracket, keys} ->
        Enum.flat_map(keys, fn key -> scan(json, key) end)

      {:deep_scan_array, array_op} ->
        collect_all_lists(json)
        |> Enum.flat_map(fn list ->
          result = eval_array(array_op, list)
          to_list_value(result)
        end)

      {:deep_scan_standalone, _} ->
        deepscan(json)

      {:array, array_op} ->
        eval_array(array_op, json)

      {:filter_expression, filters} ->
        eval_filter(json, filters, original_json)

      {:function, function} ->
        eval_function(json, function)

      {:value, [value]} ->
        value
    end
  end

  defp to_list_value(@missing), do: []
  defp to_list_value(result) when is_list(result), do: result
  defp to_list_value(result), do: [result]

  # ---------------------------------------------------------------------------
  # Key access
  # ---------------------------------------------------------------------------
  # On a list: collect values only from maps that actually have the key,
  # preserving order. Missing keys are skipped (legacy list-iteration behavior).
  defp get(json, key) when is_list(json) do
    json
    |> Enum.reduce([], fn item, acc ->
      case item do
        m when is_map(m) ->
          case Map.fetch(m, key) do
            {:ok, val} -> [val | acc]
            :error -> acc
          end

        _ ->
          acc
      end
    end)
    |> Enum.reverse()
  end

  defp get(json, key) when is_map(json) do
    Map.get(json, key, @missing)
  end

  defp get(_other, _key), do: @missing

  # Multi-key bracket child on a map: preserve bracket order, skip missing keys.
  defp multi_key_bracket(json, keys) when is_map(json) do
    keys
    |> Enum.reduce([], fn key, acc ->
      case Map.fetch(json, key) do
        {:ok, val} -> [val | acc]
        :error -> acc
      end
    end)
    |> Enum.reverse()
  end

  defp multi_key_bracket(json, keys) when is_list(json) do
    Enum.flat_map(keys, &get(json, &1))
  end

  defp multi_key_bracket(_, _), do: []

  # ---------------------------------------------------------------------------
  # Recursive descent helpers
  # ---------------------------------------------------------------------------
  defp deepscan(data) when is_map(data) do
    Enum.flat_map(data, fn {_k, v} -> [v | deepscan(v)] end)
  end

  defp deepscan(data) when is_list(data) do
    Enum.flat_map(data, fn v -> [v | deepscan(v)] end)
  end

  defp deepscan(_), do: []

  # scan: find all values for a given key at any depth, preserving order
  defp scan(map, key) when is_map(map) do
    Enum.flat_map(map, fn
      {^key, value} -> [value]
      {_k, value} when is_map(value) or is_list(value) -> scan(value, key)
      {_k, _v} -> []
    end)
  end

  defp scan(list, key) when is_list(list) do
    Enum.flat_map(list, &scan(&1, key))
  end

  defp scan(_other, _key), do: []

  # Collect every list inside the data tree. Includes the root if it's a list,
  # every list-typed map value, and every nested list inside a list.
  # Lists are NOT included twice: each list contributes once at the level it
  # appears.
  defp collect_all_lists(data) when is_list(data) do
    [data | Enum.flat_map(data, &nested_lists/1)]
  end

  defp collect_all_lists(data), do: nested_lists(data)

  defp nested_lists(data) when is_list(data) do
    [data | Enum.flat_map(data, &nested_lists/1)]
  end

  defp nested_lists(data) when is_map(data) do
    Enum.flat_map(data, fn {_k, v} -> nested_lists(v) end)
  end

  defp nested_lists(_), do: []

  # ---------------------------------------------------------------------------
  # Array operations
  # ---------------------------------------------------------------------------
  defp eval_array(array_op, json) when is_list(json) do
    case array_op do
      [array_wildcard: _] ->
        json

      # Single-index fast path: avoid the cost of building a tuple.
      [array_indices: [i]] ->
        single_index_at(json, i)

      [array_indices: indices] ->
        multi_index_at(json, indices)

      [array_slice: slice_params] ->
        eval_slice(json, slice_params)
    end
  end

  # [*] on a map → its values
  defp eval_array([array_wildcard: _], json) when is_map(json), do: Map.values(json)

  defp eval_array(_, _), do: []

  # Single-index access avoids the tuple allocation. For negative indices
  # we first need the length to normalize; for non-negative ones a forward
  # walk via Enum.at/2 is O(i) and allocation-free.
  defp single_index_at(json, i) when i >= 0 do
    case Enum.at(json, i, @no_default) do
      @no_default -> []
      val -> [val]
    end
  end

  defp single_index_at(json, i) do
    len = length(json)
    idx = len + i

    if idx >= 0 and idx < len do
      [Enum.at(json, idx)]
    else
      []
    end
  end

  # Multi-index access: O(n) length + O(n) tuple build + O(k) lookups.
  defp multi_index_at(json, indices) do
    len = length(json)
    tuple = List.to_tuple(json)

    indices
    |> Enum.reduce([], fn i, acc ->
      idx = if i >= 0, do: i, else: len + i

      if idx >= 0 and idx < len do
        [elem(tuple, idx) | acc]
      else
        acc
      end
    end)
    |> Enum.reverse()
  end

  # RFC 9535 slice semantics, with O(n) materialization (tuple-based indexing).
  defp eval_slice(json, slice_params) do
    len = length(json)
    step = Keyword.get(slice_params, :step, 1)

    if step == 0 do
      []
    else
      {start, stop} =
        if step > 0 do
          s = Keyword.get(slice_params, :begin, 0)
          e = Keyword.get(slice_params, :end, len)
          {max(normalize_index(s, len), 0), min(normalize_index(e, len), len)}
        else
          s = Keyword.get(slice_params, :begin, len - 1)
          e = Keyword.get(slice_params, :end, -len - 1)
          {min(normalize_index(s, len), len - 1), max(normalize_index(e, len), -1)}
        end

      tuple = List.to_tuple(json)

      start
      |> generate_indices(stop, step)
      |> Enum.reduce([], fn i, acc ->
        if i >= 0 and i < len, do: [elem(tuple, i) | acc], else: acc
      end)
      |> Enum.reverse()
    end
  end

  defp normalize_index(i, len) do
    if i >= 0, do: i, else: len + i
  end

  defp generate_indices(start, stop, step) when step > 0 do
    generate_indices_pos(start, stop, step, [])
  end

  defp generate_indices(start, stop, step) when step < 0 do
    generate_indices_neg(start, stop, step, [])
  end

  defp generate_indices_pos(i, stop, step, acc) when i < stop do
    generate_indices_pos(i + step, stop, step, [i | acc])
  end

  defp generate_indices_pos(_, _, _, acc), do: Enum.reverse(acc)

  defp generate_indices_neg(i, stop, step, acc) when i > stop do
    generate_indices_neg(i + step, stop, step, [i | acc])
  end

  defp generate_indices_neg(_, _, _, acc), do: Enum.reverse(acc)

  # ---------------------------------------------------------------------------
  # Filter evaluation
  # ---------------------------------------------------------------------------
  defp eval_filter(json, filters, original_json) when is_map(json) do
    eval_filter(Map.values(json), filters, original_json)
  end

  defp eval_filter(json, filters, original_json) when is_list(json) do
    Enum.filter(json, fn x ->
      Enum.all?(filters, fn filter -> truthy?(eval_term(x, filter, original_json)) end)
    end)
  end

  defp eval_filter(_, _, _), do: []

  # Truthiness used by filter selection: the missing sentinel and empty lists
  # are falsy; nil and false are falsy (Elixir default); everything else truthy.
  defp truthy?(@missing), do: false
  defp truthy?([]), do: false
  defp truthy?(nil), do: false
  defp truthy?(false), do: false
  defp truthy?(_), do: true

  # ---------------------------------------------------------------------------
  # Term / expression evaluation with proper operator precedence:
  #   arithmetic  (+, -, *, /, %)     — highest
  #   comparison  (==, <, >, in …)
  #   logical     (&&)
  #   logical     (||)                — lowest
  # ---------------------------------------------------------------------------
  defp eval_term(json, {:term, term}, original_json) do
    sequence =
      term
      |> Enum.reduce([], fn node, acc -> eval_node(acc, node, json, original_json) end)
      |> Enum.reverse()

    eval_logical_or(sequence)
  end

  # Split on || (lowest precedence)
  defp eval_logical_or(tokens) do
    tokens
    |> split_on(:||)
    |> Enum.map(&eval_logical_and/1)
    |> Enum.reduce(false, fn val, acc -> acc || val end)
  end

  # Split on &&
  defp eval_logical_and(tokens) do
    tokens
    |> split_on(:&&)
    |> Enum.map(&eval_comparison_chain/1)
    |> Enum.reduce(true, fn val, acc -> acc && val end)
  end

  # Prefix ! then arithmetic + comparison
  defp eval_comparison_chain([:! | rest]) do
    !truthy?(eval_comparison_chain(rest))
  end

  defp eval_comparison_chain(tokens) do
    tokens
    |> group_by_comparisons()
    |> Enum.map(&eval_arith/1)
    |> Enum.flat_map(fn
      [op] when op in @comparison_operators -> [op]
      rpn -> evaluate_rpn(rpn)
    end)
    |> compare()
  end

  # Split a token list on every occurrence of *sep*, returning sub-lists.
  defp split_on(list, sep) do
    {groups, current} =
      Enum.reduce(list, {[], []}, fn
        ^sep, {groups, current} -> {[Enum.reverse(current) | groups], []}
        elem, {groups, current} -> {groups, [elem | current]}
      end)

    Enum.reverse([Enum.reverse(current) | groups])
  end

  # ---------------------------------------------------------------------------
  # Comparison chain  [left, op, right, op, right …]
  # ---------------------------------------------------------------------------
  defp compare([]), do: nil
  defp compare([result]), do: result

  # `in` with a list-typed RHS
  defp compare([n1, :in, list | rest]) when is_list(list) do
    compare([n1 != @missing and Enum.member?(list, n1) | rest])
  end

  defp compare([n1, :"not in", list | rest]) when is_list(list) do
    compare([n1 != @missing and !Enum.member?(list, n1) | rest])
  end

  # Missing-vs-missing equality (Nothing == Nothing is true per RFC 9535)
  defp compare([@missing, op, @missing | rest]) when op in [:==, :===] do
    compare([true | rest])
  end

  defp compare([@missing, op, @missing | rest]) when op == :!= do
    compare([false | rest])
  end

  defp compare([@missing, op, _ | rest]) when op in [:==, :===] do
    compare([false | rest])
  end

  defp compare([_, op, @missing | rest]) when op in [:==, :===] do
    compare([false | rest])
  end

  defp compare([@missing, op, _ | rest]) when op == :!= do
    compare([true | rest])
  end

  defp compare([_, op, @missing | rest]) when op == :!= do
    compare([true | rest])
  end

  # Ordering comparisons involving missing → always false
  defp compare([@missing, op, _ | rest]) when op in @ordering_operators do
    compare([false | rest])
  end

  defp compare([_, op, @missing | rest]) when op in @ordering_operators do
    compare([false | rest])
  end

  # Ordering comparisons involving nil → always false (treat as Nothing for
  # backward compatibility — nil from a list-iteration get/2).
  defp compare([nil, op, _ | rest]) when op in @ordering_operators do
    compare([false | rest])
  end

  defp compare([_, op, nil | rest]) when op in @ordering_operators do
    compare([false | rest])
  end

  # Ordering operators only meaningful between two numbers OR two strings.
  defp compare([n1, op, n2 | rest]) when op in @ordering_operators do
    cond do
      is_number(n1) and is_number(n2) ->
        compare([apply(Kernel, op, [n1, n2]) | rest])

      is_binary(n1) and is_binary(n2) ->
        compare([apply(Kernel, op, [n1, n2]) | rest])

      true ->
        compare([false | rest])
    end
  end

  defp compare([n1, op, n2 | rest]) when op in [:==, :!=, :===] do
    compare([apply(Kernel, op, [n1, n2]) | rest])
  end

  defp compare([n1, op, n2 | rest]) do
    compare([apply(Kernel, op, [n1, n2]) | rest])
  end

  # ---------------------------------------------------------------------------
  # Node evaluation (builds the flat token sequence)
  # ---------------------------------------------------------------------------
  defp eval_node(sequence, node, json, original_json) do
    case node do
      # bare ! in the token stream (prefix negation without grouping)
      :! ->
        [:! | sequence]

      # negated grouping  !( … )
      {:grouping, [:!, group]} ->
        [!truthy?(eval_term(json, group, original_json)) | sequence]

      # plain grouping  ( … )
      {:grouping, [group]} ->
        [eval_term(json, group, original_json) | sequence]

      # scalar / string / boolean / null literal
      {:operand, {:value, value}} ->
        [value | sequence]

      # list literal  [2, 3]  — values are tagged {:value, v} by the parser
      {:operand, {:list_value, items}} ->
        values = Enum.map(items, fn {:value, v} -> v end)
        [values | sequence]

      # @.path  or just @
      {:operand, {:current_context, [{:current, _} | rest]}} ->
        [unwrap_nodelist(evaluate1(rest, json, original_json), rest) | sequence]

      # $.path  (root reference inside a filter)
      {:operand, {:root_key, rest}} ->
        [unwrap_nodelist(evaluate1(rest, original_json, original_json), rest) | sequence]

      # prefix function call:  length(@), count(@.x), match(@.s, "regex")
      {:operand, {:prefix_function, {:prefix_call, name, args}}} ->
        evaluated = Enum.map(args, &eval_prefix_arg(&1, json, original_json))
        [eval_prefix_function(name, evaluated) | sequence]

      {:operator, operator} ->
        [operator | sequence]
    end
  end

  defp eval_prefix_arg({:current_context, [{:current, _} | rest]}, json, original_json) do
    unwrap_nodelist(evaluate1(rest, json, original_json), rest)
  end

  defp eval_prefix_arg({:root_key, rest}, _json, original_json) do
    unwrap_nodelist(evaluate1(rest, original_json, original_json), rest)
  end

  defp eval_prefix_arg({:value, value}, _json, _original_json), do: value

  # ---------------------------------------------------------------------------
  # Single-index array access in a filter operand path produces a one-element
  # nodelist; unwrap it to a scalar so comparisons work naturally.
  # Other list results (wildcards, slices, filters) stay as lists.
  # ---------------------------------------------------------------------------
  defp unwrap_nodelist(@missing, _path), do: @missing

  defp unwrap_nodelist([value], path) do
    case List.last(path) do
      {:array, [array_indices: [_]]} -> value
      _ -> [value]
    end
  end

  defp unwrap_nodelist(result, _path), do: result

  # ---------------------------------------------------------------------------
  # Built-in functions (postfix on a path, e.g. $.items.length())
  # ---------------------------------------------------------------------------
  defp eval_function(json, function) do
    apply_function(function, json)
  end

  defp apply_function(:length, json) when is_list(json), do: length(json)
  defp apply_function(:length, json) when is_map(json), do: map_size(json)
  defp apply_function(:length, json) when is_binary(json), do: String.length(json)
  defp apply_function(:length, _), do: nil

  defp apply_function(:count, json) when is_list(json), do: length(json)
  defp apply_function(:count, json) when is_map(json), do: map_size(json)
  defp apply_function(:count, _), do: nil

  defp apply_function(:sum, json) when is_list(json) do
    nums = Enum.filter(json, &is_number/1)
    if nums == [], do: 0, else: Enum.sum(nums)
  end

  defp apply_function(:sum, _), do: nil

  defp apply_function(:min, json) when is_list(json) do
    nums = Enum.filter(json, &is_number/1)

    case nums do
      [] -> nil
      _ -> Enum.min(nums)
    end
  end

  defp apply_function(:min, _), do: nil

  defp apply_function(:max, json) when is_list(json) do
    nums = Enum.filter(json, &is_number/1)

    case nums do
      [] -> nil
      _ -> Enum.max(nums)
    end
  end

  defp apply_function(:max, _), do: nil

  defp apply_function(:avg, json) when is_list(json) do
    nums = Enum.filter(json, &is_number/1)

    case nums do
      [] -> nil
      _ -> Enum.sum(nums) / length(nums)
    end
  end

  defp apply_function(:avg, _), do: nil

  defp apply_function(:concat, json) when is_list(json) do
    json |> Enum.map(&to_concat_str/1) |> Enum.join()
  end

  defp apply_function(:concat, _), do: nil

  defp to_concat_str(nil), do: ""
  defp to_concat_str(s) when is_binary(s), do: s
  defp to_concat_str(n) when is_number(n), do: to_string(n)
  defp to_concat_str(true), do: "true"
  defp to_concat_str(false), do: "false"
  defp to_concat_str(other), do: inspect(other)

  # ---------------------------------------------------------------------------
  # Prefix functions invoked inside filter expressions.
  # ---------------------------------------------------------------------------
  defp eval_prefix_function(name, [arg])
       when name in [:length, :count, :min, :max, :sum, :avg, :concat] do
    apply_function(name, prefix_arg_to_value(arg))
  end

  defp eval_prefix_function(:match, [s, pattern]) when is_binary(s) and is_binary(pattern) do
    case Regex.compile("\\A(?:" <> pattern <> ")\\z") do
      {:ok, re} -> Regex.match?(re, s)
      {:error, _} -> false
    end
  end

  defp eval_prefix_function(:search, [s, pattern]) when is_binary(s) and is_binary(pattern) do
    case Regex.compile(pattern) do
      {:ok, re} -> Regex.match?(re, s)
      {:error, _} -> false
    end
  end

  defp eval_prefix_function(:match, _), do: false
  defp eval_prefix_function(:search, _), do: false
  defp eval_prefix_function(_, _), do: nil

  # For prefix function arguments, treat the missing sentinel as nil so callers
  # like length(@.foo) on a missing path return nil rather than crashing.
  defp prefix_arg_to_value(@missing), do: nil
  defp prefix_arg_to_value(other), do: other

  # ---------------------------------------------------------------------------
  # Arithmetic helpers (shunting-yard → RPN → evaluate)
  # ---------------------------------------------------------------------------
  # Split the flat token list into segments separated by comparison operators.
  defp group_by_comparisons(list) do
    {groups, current} =
      Enum.reduce(list, {[], []}, fn
        elem, {groups, current} when elem in @comparison_operators ->
          [Enum.reverse(current) | groups]
          |> then(fn acc -> {[[elem] | acc], []} end)

        elem, {groups, current} ->
          {groups, [elem | current]}
      end)

    Enum.reverse([Enum.reverse(current) | groups])
  end

  # Shunting-yard: infix arithmetic → RPN.
  defp eval_arith(list) do
    precedence = %{:* => 2, :/ => 2, :% => 2, :+ => 1, :- => 1}

    {output, stack} =
      Enum.reduce(list, {[], []}, fn
        op, {output, stack} when op in @arithmetic_operators ->
          {to_pop, new_stack} =
            Enum.split_while(stack, fn s ->
              Map.get(precedence, s, 0) >= Map.get(precedence, op)
            end)

          {Enum.reverse(to_pop) ++ output, [op | new_stack]}

        x, {output, stack} ->
          {[x | output], stack}
      end)

    Enum.reverse(output) ++ stack
  end

  # RPN evaluator: only arithmetic operators trigger apply; everything else
  # (numbers, strings, booleans, nil, lists) is a value.
  # Division by zero and non-numeric operands propagate the missing sentinel
  # so downstream comparisons safely return false.
  defp evaluate_rpn(rpn) do
    Enum.reduce(rpn, [], fn
      op, [n2, n1 | rest] when op in @arithmetic_operators ->
        [arith(op, n1, n2) | rest]

      val, stack ->
        [val | stack]
    end)
  end

  defp arith(_, @missing, _), do: @missing
  defp arith(_, _, @missing), do: @missing
  defp arith(:/, _, 0), do: @missing
  defp arith(:/, _, 0.0), do: @missing
  defp arith(:%, _, 0), do: @missing
  defp arith(:%, _, 0.0), do: @missing

  defp arith(:%, n1, n2) when is_integer(n1) and is_integer(n2) do
    rem(n1, n2)
  end

  defp arith(:%, n1, n2) when is_number(n1) and is_number(n2) do
    :math.fmod(n1 * 1.0, n2 * 1.0)
  end

  defp arith(op, n1, n2) when op in @arithmetic_operators and is_number(n1) and is_number(n2) do
    apply(Kernel, op, [n1, n2])
  end

  defp arith(_, _, _), do: @missing
end