guesswork

defmodule Guesswork.Ast.Or do
  @moduledoc """
  Simple disjunction.
  """

  alias Guesswork.Ast.Term
  alias Guesswork.Ast.Statement

  defstruct [:args]

  @type t() :: %__MODULE__{args: [Statement.t()]}

  @spec new([Statement.t()]) :: t()
  def new(args), do: %__MODULE__{args: args}

  defimpl Term do
    alias Guesswork.Ast.Or

    def concrete?(%{args: args}), do: Enum.all?(args, &Term.concrete?/1)

    def get_variables(%Or{args: args}) do
      Enum.flat_map(args, &Term.get_variables/1)
    end
  end

  defimpl Statement do
    alias Guesswork.Ast.Or

    def find_statements(%Or{args: args} = statement, fun) do
      rest = Enum.flat_map(args, &Statement.find_statements(&1, fun))

      if fun.(statement) do
        [statement | rest]
      else
        rest
      end
    end

    def simplify(%{args: []}), do: nil
    def simplify(%{args: [arg]}), do: Statement.simplify(arg)

    def simplify(%{args: args}) do
      case args
           |> Enum.map(&Statement.simplify/1)
           |> Enum.filter(&Kernel.!=(nil, &1))
           |> Enum.flat_map(&raise_args/1) do
        [] -> nil
        args -> Or.new(args)
      end
    end

    defp raise_args(%Or{args: args}), do: args
    defp raise_args(val), do: [val]

    def substitute(%Or{args: args}, env),
      do: %Or{args: Enum.map(args, &Statement.substitute(&1, env))}

    def resolve(%Or{args: args}, opts) do
      {calls, streams} =
        args
        |> Enum.map(&Statement.resolve(&1, opts))
        |> Enum.split_with(&is_function(&1, 1))

      call_stream = calls |> Enum.reverse() |> Stream.map(&apply(&1, [%{}]))

      Stream.concat(streams, call_stream) |> Stream.flat_map(& &1)
    end
  end

  defimpl String.Chars do
    alias Guesswork.Ast.Or

    def to_string(%Or{args: args}),
      do: "l_or(#{Enum.map_join(args, ", ", &String.Chars.to_string/1)})"
  end
end