guesswork

defmodule Guesswork.Query do
  @moduledoc """
  Represents an un-resolved query.
  """

  defstruct [:computation, :id, :label]

  alias Guesswork.Answer
  alias Guesswork.Answer.Result
  alias Guesswork.Ast.Fact
  alias Guesswork.Ast.Statement
  alias Guesswork.Ast.Statement.Opts
  alias Guesswork.Telemetry.EventHandler

  @type t() :: %__MODULE__{
          id: String.t(),
          computation: Enumerable.t(Answer.t()),
          label: boolean()
        }

  @opts_schema_list [
                      label: [
                        type: {:or, [:boolean, {:list, :string}]},
                        default: false,
                        doc: """
                        Traverses plural entities looking for a specific answer sets.
                        If a list of strings is provided only the supplied variables are traversed.
                        """
                      ]
                    ] ++ Opts.schema_list()
  @opts_schema NimbleOptions.new!(@opts_schema_list)
  @type opts_list() :: [unquote(NimbleOptions.option_typespec(@opts_schema))]

  @doc """
  Builds a new query.

  Available Options:

  #{NimbleOptions.docs(@opts_schema)}
  """
  @spec new(Statement.t(), opts_list()) :: t()
  def new(statement, opts \\ [])

  def new(%Fact{} = statement, opts) do
    label = Keyword.get(opts, :label, false)
    %Opts{query_id: query_id} = opts = Opts.new(opts)

    computation =
      statement
      |> Statement.simplify()
      |> Statement.resolve(opts)
      |> apply([%{}])
      |> handle_labeling(label)
      |> Stream.map(&Answer.remove_resolvable_constraints/1)
      |> Stream.map(&Answer.remove_internal_variables/1)

    %__MODULE__{computation: computation, id: query_id, label: label}
  end

  def new(statement, opts) do
    label = Keyword.get(opts, :label, false)
    %Opts{query_id: query_id} = opts = Opts.new(opts)

    computation =
      statement
      |> Statement.simplify()
      |> Statement.resolve(opts)
      |> handle_labeling(label)
      |> Stream.map(&Answer.remove_resolvable_constraints/1)
      |> Stream.map(&Answer.remove_internal_variables/1)

    %__MODULE__{computation: computation, id: query_id, label: label}
  end

  defp handle_labeling(stream, true), do: Stream.flat_map(stream, &Answer.traverse_entities/1)
  defp handle_labeling(stream, false), do: stream

  defp handle_labeling(stream, vars),
    do: Stream.flat_map(stream, &Answer.traverse_entities(&1, MapSet.new(vars)))

  def opts_schema() do
    @opts_schema
  end

  @event_prefix EventHandler.resolve_span_prefix()

  @doc """
  Takes a query and attempts to resolve `n` answers from it.
  """
  @spec run(t(), integer()) :: Result.t()
  def run(%__MODULE__{computation: computation, id: id}, n \\ 1) do
    metadata = %{query_id: id, run_id: UUID.uuid4(), n: n}

    answers =
      :telemetry.span(@event_prefix, metadata, fn ->
        result = Enum.take(computation, n)
        {result, metadata}
      end)

    %Result{answer_sets: Enum.map(answers, &Answer.flatten/1)}
  end
end