defmodule AshSql.Join do
  @moduledoc false
  import Ecto.Query, only: [from: 2, subquery: 1]

  alias Ash.Query.{Not, Ref}

  @known_inner_join_operators [
    Ash.Query.Operator.Eq,
    Ash.Query.Operator.GreaterThan,
    Ash.Query.Operator.GreaterThanOrEqual,
    Ash.Query.Operator.In,
    Ash.Query.Operator.LessThanOrEqual,
    Ash.Query.Operator.LessThan,
    Ash.Query.Operator.NotEq
  ]

  @known_inner_join_functions [
    Ash.Query.Function.Ago,
    Ash.Query.Function.Contains,
    Ash.Query.Function.At,
    Ash.Query.Function.DateAdd,
    Ash.Query.Function.DateTimeAdd,
    Ash.Query.Function.FromNow,
    Ash.Query.Function.GetPath,
    Ash.Query.Function.Length,
    Ash.Query.Function.Minus,
    Ash.Query.Function.Round,
    Ash.Query.Function.StringDowncase,
    Ash.Query.Function.StringJoin,
    Ash.Query.Function.StringLength,
    Ash.Query.Function.StringSplit,
    Ash.Query.Function.StringTrim
  ]

  def join_all_relationships(
        query,
        filter,
        opts \\ [],
        relationship_paths \\ nil,
        path \\ [],
        source \\ nil,
        sort? \\ true,
        join_filters \\ nil,
        parent_bindings \\ nil,
        no_inner_join? \\ false
      )

  # simple optimization for common cases
  def join_all_relationships(
        query,
        filter,
        _opts,
        relationship_paths,
        _path,
        _source,
        _sort?,
        _join_filters,
        _parent_bindings,
        _no_inner_join?
      )
      when is_nil(relationship_paths) and filter in [nil, true, false] do
    {:ok, query}
  end

  def join_all_relationships(
        query,
        filter,
        opts,
        relationship_paths,
        path,
        source,
        sort?,
        join_filters,
        parent_query,
        no_inner_join?
      ) do
    case join_parent_paths(query, filter, relationship_paths) do
      {:ok, query} ->
        relationship_paths =
          relationship_paths ||
            filter
            |> Ash.Filter.relationship_paths()
            |> to_joins(filter, query.__ash_bindings__.resource)

        Enum.reduce_while(relationship_paths, {:ok, query}, fn
          {_join_type, []}, {:ok, query} ->
            {:cont, {:ok, query}}

          {join_type, [relationship | rest_rels]}, {:ok, query} ->
            join_type =
              if no_inner_join? do
                :left
              else
                join_type
              end

            source = source || relationship.source

            current_path = path ++ [relationship]

            current_join_type = join_type

            look_for_join_types =
              case join_type do
                :left ->
                  [:left, :inner]

                :inner ->
                  [:left, :inner]

                other ->
                  [other]
              end

            binding =
              get_binding(source, Enum.map(current_path, & &1.name), query, look_for_join_types)

            # We can't reuse joins if we're adding filters/have a separate parent binding
            if is_nil(join_filters) && is_nil(parent_query) && binding do
              case join_all_relationships(
                     query,
                     filter,
                     opts,
                     [{join_type, rest_rels}],
                     current_path,
                     source,
                     sort?
                   ) do
                {:ok, query} ->
                  {:cont, {:ok, query}}

                {:error, error} ->
                  {:halt, {:error, error}}
              end
            else
              case join_relationship(
                     query,
                     relationship,
                     Enum.map(path, & &1.name),
                     current_join_type,
                     source,
                     filter,
                     sort?,
                     join_filters[Enum.map(current_path, & &1.name)]
                   ) do
                {:ok, joined_query} ->
                  joined_query_with_distinct = add_distinct(relationship, join_type, joined_query)

                  case join_all_relationships(
                         joined_query_with_distinct,
                         filter,
                         opts,
                         [{join_type, rest_rels}],
                         current_path,
                         source,
                         sort?,
                         join_filters,
                         Map.update!(joined_query, :__ash_bindings__, fn ash_bindings ->
                           Map.put(
                             ash_bindings,
                             :refs_at_path,
                             Enum.map(path, & &1.name) ++ [relationship.name]
                           )
                         end)
                       ) do
                    {:ok, query} ->
                      {:cont, {:ok, query}}

                    {:error, error} ->
                      {:halt, {:error, error}}
                  end

                {:error, error} ->
                  {:halt, {:error, error}}
              end
            end
        end)

      {:error, error} ->
        {:error, error}
    end
  end

  defp join_parent_paths(query, filter, nil) do
    parent_exprs =
      Ash.Filter.flat_map(filter, fn
        %Ash.Query.Parent{expr: expr} ->
          [expr]

        _ ->
          []
      end)

    case query.__ash_bindings__[:lateral_join_source_query] do
      nil ->
        {:ok, query}

      lateral_join_source_query ->
        case join_all_relationships(lateral_join_source_query, parent_exprs) do
          {:ok, lateral_join_source_query} ->
            {:ok,
             put_in(query.__ash_bindings__.lateral_join_source_query, lateral_join_source_query)
             |> Map.update!(:__ash_bindings__, fn bindings ->
               Map.put(
                 bindings,
                 :parent_bindings,
                 Map.put(lateral_join_source_query.__ash_bindings__, :parent?, true)
               )
             end)}

          {:error, error} ->
            {:error, error}
        end
    end
  end

  defp join_parent_paths(query, _filter, _relationship_paths) do
    {:ok, query}
  end

  defp to_joins(paths, filter, resource) do
    paths
    |> Enum.reject(&(&1 == []))
    |> Enum.map(fn path ->
      if can_inner_join?(path, filter) do
        {:inner,
         relationship_path_to_relationships(
           resource,
           path
         )}
      else
        {:left,
         relationship_path_to_relationships(
           resource,
           path
         )}
      end
    end)
  end

  def relationship_path_to_relationships(resource, path, acc \\ [])
  def relationship_path_to_relationships(_resource, [], acc), do: Enum.reverse(acc)

  def relationship_path_to_relationships(resource, [name | rest], acc) do
    relationship = Ash.Resource.Info.relationship(resource, name)

    if !relationship do
      raise "no such relationship #{inspect(resource)}.#{name}"
    end

    relationship_path_to_relationships(relationship.destination, rest, [relationship | acc])
  end

  def related_subquery(
        relationship,
        root_query,
        opts \\ []
      ) do
    on_parent_expr = Keyword.get(opts, :on_parent_expr, & &1)
    on_subquery = Keyword.get(opts, :on_subquery, & &1)
    filter = Keyword.get(opts, :filter, nil)
    filter_subquery? = Keyword.get(opts, :filter_subquery?, false)

    with {:ok, query} <- related_query(relationship, root_query, opts) do
      has_parent_expr? =
        !!query.__ash_bindings__.context[:data_layer][:has_parent_expr?] ||
          not is_nil(query.limit)

      query =
        if has_parent_expr? do
          on_parent_expr.(query)
        else
          query
        end

      query = on_subquery.(query)

      query = limit_from_many(query, relationship, filter, filter_subquery?, opts)

      query =
        if opts[:return_subquery?] do
          subquery(query)
        else
          if Enum.empty?(query.joins) && Enum.empty?(query.order_bys) && Enum.empty?(query.wheres) do
            query
          else
            from(row in subquery(query), as: ^(opts[:start_bindings_at] || 0))
            |> AshSql.Bindings.default_bindings(
              relationship.destination,
              query.__ash_bindings__.sql_behaviour
            )
            |> AshSql.Bindings.merge_expr_accumulator(
              query.__ash_bindings__.expression_accumulator
            )
            |> Map.update!(
              :__ash_bindings__,
              fn bindings ->
                bindings
                |> Map.put(:current, query.__ash_bindings__.current)
                |> put_in([:context, :data_layer], %{
                  has_parent_expr?: has_parent_expr?
                })
              end
            )
          end
        end

      {:ok, query}
    end
  end

  defp related_query(relationship, query, opts) do
    sort? = Keyword.get(opts, :sort?, false)
    filter = Keyword.get(opts, :filter, nil)
    filter_subquery? = Keyword.get(opts, :filter_subquery?, false)
    parent_resources = Keyword.get(opts, :parent_stack, [relationship.source])

    read_action =
      relationship.read_action ||
        Ash.Resource.Info.primary_action!(relationship.destination, :read).name

    context = query.__ash_bindings__.context

    relationship.destination
    |> Ash.Query.new()
    |> Ash.Query.set_context(context)
    |> Ash.Query.set_context(%{data_layer: %{in_group?: !!opts[:in_group?]}})
    |> Ash.Query.set_context(%{
      data_layer: %{table: nil, start_bindings_at: opts[:start_bindings_at] || 0}
    })
    |> Ash.Query.set_context(relationship.context)
    |> Ash.Query.do_filter(relationship.filter, parent_stack: parent_resources)
    |> then(fn query ->
      if Map.get(relationship, :from_many?) && filter_subquery? do
        query
      else
        Ash.Query.do_filter(query, filter, parent_stack: parent_resources)
      end
    end)
    |> Ash.Query.do_filter(opts[:apply_filter], parent_stack: parent_resources)
    |> Ash.Query.for_read(read_action, %{},
      actor: context[:private][:actor],
      tenant: context[:private][:tenant]
    )
    |> Ash.Query.set_tenant(Map.get(query, :__tenant__))
    |> Ash.Query.unset([:sort, :distinct, :select, :limit, :offset])
    |> hydrate_refs(context[:private][:actor])
    |> then(fn query ->
      if sort? do
        Ash.Query.sort(query, relationship.sort)
      else
        Ash.Query.unset(query, :sort)
      end
    end)
    |> set_has_parent_expr_context(relationship)
    |> case do
      %{valid?: true} = related_query ->
        Ash.Query.data_layer_query(
          Ash.Query.set_context(related_query, %{
            data_layer: %{
              parent_bindings:
                Map.put(query.__ash_bindings__, :refs_at_path, List.wrap(opts[:refs_at_path]))
            }
          })
        )
        |> case do
          {:ok, ecto_query} ->
            {:ok,
             ecto_query
             |> set_join_prefix(query, relationship.destination)
             |> Ecto.Query.exclude(:select)}

          {:error, error} ->
            {:error, error}
        end

      %{errors: errors} ->
        {:error, errors}
    end
  end

  defp hydrate_refs(query, actor) do
    query.filter
    |> Ash.Expr.fill_template(actor, %{}, query.context)
    |> Ash.Filter.hydrate_refs(%{resource: query.resource})
    |> case do
      {:ok, result} -> %{query | filter: result}
      {:error, error} -> Ash.Query.add_error(query, error)
    end
  end

  defp limit_from_many(
         query,
         %{from_many?: true, destination: destination},
         filter,
         filter_subquery?,
         opts
       ) do
    if filter_subquery? do
      query =
        from(row in Ecto.Query.subquery(from(row in query, limit: 1)),
          as: ^query.__ash_bindings__.root_binding
        )
        |> Map.put(:__ash_bindings__, query.__ash_bindings__)
        |> AshSql.Bindings.default_bindings(
          destination,
          query.__ash_bindings__.sql_behaviour
        )

      {:ok, query} = AshSql.Filter.filter(query, filter, query.__ash_bindings__.resource)

      if opts[:select_star?] do
        from(row in Ecto.Query.exclude(query, :select), select: 1)
      else
        query
      end
    else
      if opts[:select_star?] do
        from(row in Ecto.Query.exclude(query, :select), select: 1)
      else
        query
      end
    end
  end

  defp limit_from_many(query, _, _, _, opts) do
    if opts[:select_star?] do
      from(row in Ecto.Query.exclude(query, :select), select: 1)
    else
      query
    end
  end

  defp set_has_parent_expr_context(query, relationship) do
    has_parent_expr? =
      Ash.Actions.Read.Relationships.has_parent_expr?(%{
        relationship
        | filter: query.filter,
          sort: query.sort
      })

    Ash.Query.set_context(query, %{data_layer: %{has_parent_expr?: has_parent_expr?}})
  end

  def set_join_prefix(join_query, query, resource) do
    if Ash.Resource.Info.multitenancy_strategy(resource) == :context do
      %{
        join_query
        | prefix:
            Map.get(query, :__tenant__) ||
              query.__ash_bindings__.sql_behaviour.schema(resource) ||
              query.prefix ||
              query.__ash_bindings__.sql_behaviour.repo(resource, :mutate).config()[
                :default_prefix
              ]
      }
    else
      %{
        join_query
        | prefix:
            query.__ash_bindings__.sql_behaviour.schema(resource) ||
              query.__ash_bindings__.sql_behaviour.repo(resource, :mutate).config()[
                :default_prefix
              ]
      }
    end
  end

  defp can_inner_join?(path, expr) do
    expr
    |> AshSql.Expr.split_statements(:and)
    |> Enum.any?(&known_inner_join_predicate_for_path_in_all_branches?(path, &1))
  end

  defp known_inner_join_predicate_for_path_in_all_branches?(path, expr) do
    expr
    |> AshSql.Expr.split_statements(:or)
    |> case do
      [expr] ->
        case AshSql.Expr.split_statements(expr, :and) do
          [expr] ->
            known_predicates_only_containing?(path, expr)

          many ->
            Enum.any?(many, &known_inner_join_predicate_for_path_in_all_branches?(path, &1))
        end

      branches ->
        Enum.all?(branches, &can_inner_join?(path, &1))
    end
  end

  defp known_predicates_only_containing?(path, %Not{expression: expression}) do
    known_predicates_only_containing?(path, expression)
  end

  defp known_predicates_only_containing?(path, %struct{
         __operator__?: true,
         left: left,
         right: right
       })
       when struct in @known_inner_join_operators do
    Enum.any?([left, right], &known_predicates_only_containing?(path, &1))
  end

  defp known_predicates_only_containing?(path, %struct{
         __function__?: true,
         arguments: arguments
       })
       when struct in @known_inner_join_functions do
    Enum.any?(arguments, &known_predicates_only_containing?(path, &1))
  end

  defp known_predicates_only_containing?(path, %Ref{relationship_path: path}) do
    true
  end

  defp known_predicates_only_containing?(_, _), do: false

  @doc false
  def get_binding(resource, candidate_path, %{__ash_bindings__: _} = query, types) do
    types = List.wrap(types)

    Enum.find_value(query.__ash_bindings__.bindings, fn
      {binding, %{path: path, source: source, type: type}} ->
        if type in types &&
             Ash.SatSolver.synonymous_relationship_paths?(resource, path, candidate_path, source) do
          binding
        end

      _ ->
        nil
    end)
  end

  def get_binding(_, _, _, _), do: nil

  defp add_distinct(relationship, _join_type, joined_query) do
    if !(joined_query.__ash_bindings__.in_group? ||
           joined_query.__ash_bindings__.context[:data_layer][:in_group?]) &&
         (relationship.cardinality == :many || Map.get(relationship, :from_many?)) &&
         !joined_query.distinct do
      pkey = Ash.Resource.Info.primary_key(joined_query.__ash_bindings__.resource)

      if joined_query.__ash_bindings__.sql_behaviour.multicolumn_distinct?() do
        from(row in joined_query,
          distinct: ^pkey
        )
      else
        from(row in joined_query, distinct: true)
      end
    else
      joined_query
    end
  end

  defp join_relationship(
         query,
         %{manual: {module, opts}} = relationship,
         path,
         kind,
         source,
         filter,
         sort?,
         apply_filter
       ) do
    full_path = path ++ [relationship.name]
    initial_ash_bindings = query.__ash_bindings__

    binding_data = %{type: kind, path: full_path, source: source}

    query = AshSql.Bindings.add_binding(query, binding_data)

    used_aggregates = Ash.Filter.used_aggregates(filter, full_path)

    with {:ok, relationship_destination} <-
           related_subquery(relationship, query,
             sort?: sort?,
             apply_filter?: apply_filter,
             refs_at_path: path
           ) do
      binding_kinds =
        case kind do
          :left ->
            [:left, :inner]

          :inner ->
            [:left, :inner]

          other ->
            [other]
        end

      current_binding =
        Enum.find_value(
          initial_ash_bindings.bindings,
          initial_ash_bindings.root_binding,
          fn {binding, data} ->
            if data.type in binding_kinds && data.path == path do
              binding
            end
          end
        )

      case apply(module, query.__ash_bindings__.sql_behaviour.manual_relationship_function(), [
             query,
             opts,
             current_binding,
             initial_ash_bindings.current,
             kind,
             relationship_destination
           ]) do
        {:ok, query} ->
          AshSql.Aggregate.add_aggregates(
            query,
            used_aggregates,
            relationship.destination,
            false,
            initial_ash_bindings.current,
            {query.__ash_bindings__.resource, full_path}
          )

        {:error, query} ->
          {:error, query}
      end
    end
  rescue
    e in UndefinedFunctionError ->
      if e.function == query.__ash_bindings__.sql_behaviour.manual_relationship_function() do
        reraise """
                Cannot join to a manual relationship #{inspect(module)} that does not implement the `#{query.__ash_bindings__.sql_behaviour.manual_relationship_behaviour()}` behaviour.
                """,
                __STACKTRACE__
      else
        reraise e, __STACKTRACE__
      end
  end

  defp join_relationship(
         query,
         %{type: :many_to_many} = relationship,
         path,
         kind,
         source,
         filter,
         sort?,
         apply_filter
       ) do
    join_relationship =
      Ash.Resource.Info.relationship(relationship.source, relationship.join_relationship)

    join_path = path ++ [join_relationship.name]

    full_path = path ++ [relationship.name]

    initial_ash_bindings = query.__ash_bindings__

    binding_data = %{type: kind, path: full_path, source: source}

    used_aggregates = Ash.Filter.used_aggregates(filter, full_path)

    query =
      query
      |> AshSql.Bindings.add_binding(%{
        path: join_path,
        type: :left,
        source: source
      })
      |> AshSql.Bindings.add_binding(binding_data)

    with {:ok, relationship_through} <- related_subquery(join_relationship, query),
         {:ok, relationship_destination} <-
           related_subquery(relationship, query,
             sort?: sort?,
             apply_filter?: apply_filter,
             refs_at_path: path
           ) do
      {relationship_destination, dest_acc} =
        maybe_apply_filter(relationship_destination, query, query.__ash_bindings__, apply_filter)

      query =
        query
        |> AshSql.Bindings.merge_expr_accumulator(dest_acc)

      binding_kinds =
        case kind do
          :left ->
            [:left, :inner]

          :inner ->
            [:left, :inner]

          other ->
            [other]
        end

      current_binding =
        Enum.find_value(
          initial_ash_bindings.bindings,
          initial_ash_bindings.root_binding,
          fn {binding, data} ->
            if data.type in binding_kinds && data.path == path do
              binding
            end
          end
        )

      query =
        case kind do
          :inner ->
            from(_ in query,
              join: through in ^relationship_through,
              as: ^initial_ash_bindings.current,
              on:
                field(as(^current_binding), ^relationship.source_attribute) ==
                  field(through, ^relationship.source_attribute_on_join_resource),
              join: destination in ^relationship_destination,
              as: ^(initial_ash_bindings.current + 1),
              on:
                field(destination, ^relationship.destination_attribute) ==
                  field(through, ^relationship.destination_attribute_on_join_resource)
            )

          _ ->
            from(_ in query,
              left_join: through in ^relationship_through,
              as: ^initial_ash_bindings.current,
              on:
                field(as(^current_binding), ^relationship.source_attribute) ==
                  field(through, ^relationship.source_attribute_on_join_resource),
              left_join: destination in ^relationship_destination,
              as: ^(initial_ash_bindings.current + 1),
              on:
                field(destination, ^relationship.destination_attribute) ==
                  field(through, ^relationship.destination_attribute_on_join_resource)
            )
        end

      AshSql.Aggregate.add_aggregates(
        query,
        used_aggregates,
        relationship.destination,
        false,
        initial_ash_bindings.current,
        {query.__ash_bindings__.resource, full_path}
      )
    end
  end

  defp join_relationship(
         query,
         relationship,
         path,
         kind,
         source,
         filter,
         sort?,
         apply_filter
       ) do
    full_path = path ++ [relationship.name]
    initial_ash_bindings = query.__ash_bindings__

    binding_data = %{type: kind, path: full_path, source: source}

    query = AshSql.Bindings.add_binding(query, binding_data)

    used_aggregates = Ash.Filter.used_aggregates(filter, full_path)

    binding_kinds =
      case kind do
        :left ->
          [:left, :inner]

        :inner ->
          [:left, :inner]

        other ->
          [other]
      end

    current_binding =
      Enum.find_value(
        initial_ash_bindings.bindings,
        initial_ash_bindings.root_binding,
        fn {binding, data} ->
          if data.type in binding_kinds && data.path == path do
            binding
          end
        end
      )

    case related_subquery(relationship, query,
           sort?: sort?,
           apply_filter: apply_filter,
           refs_at_path: path,
           filter_subquery?: true,
           sort?: Map.get(relationship, :from_many?),
           on_subquery: fn subquery ->
             if !Map.get(relationship, :from_many?) || Map.get(relationship, :no_attributes?) do
               subquery
             else
               source_ref =
                 AshSql.Expr.ref_binding(
                   %Ref{
                     attribute:
                       Ash.Resource.Info.attribute(
                         query.__ash_bindings__.resource,
                         relationship.source_attribute
                       ),
                     resource: query.__ash_bindings__.resource,
                     relationship_path: []
                   },
                   query.__ash_bindings__
                 )

               Ecto.Query.from(destination in subquery,
                 where:
                   field(parent_as(^source_ref), ^relationship.source_attribute) ==
                     field(destination, ^relationship.destination_attribute)
               )
             end
           end,
           on_parent_expr: fn subquery ->
             if Map.get(relationship, :no_attributes?) do
               subquery
             else
               from(row in subquery,
                 where:
                   field(parent_as(^current_binding), ^relationship.source_attribute) ==
                     field(
                       row,
                       ^relationship.destination_attribute
                     )
               )
             end
           end
         ) do
      {:error, error} ->
        {:error, error}

      {:ok, relationship_destination} ->
        query =
          case {kind, Map.get(relationship, :no_attributes?, false),
                relationship_destination.__ash_bindings__.context[:data_layer][
                  :has_parent_expr?
                ] || Map.get(relationship, :from_many?, false)} do
            {:inner, true, false} ->
              from(_ in query,
                join: destination in ^relationship_destination,
                as: ^initial_ash_bindings.current,
                on: true
              )

            {:inner, true, true} ->
              from(_ in query,
                inner_lateral_join: destination in ^relationship_destination,
                as: ^initial_ash_bindings.current,
                on: true
              )

            {:inner, false, false} ->
              from(_ in query,
                join: destination in ^relationship_destination,
                as: ^initial_ash_bindings.current,
                on:
                  field(as(^current_binding), ^relationship.source_attribute) ==
                    field(
                      destination,
                      ^relationship.destination_attribute
                    )
              )

            {:inner, false, true} ->
              from(_ in query,
                inner_lateral_join: destination in ^relationship_destination,
                as: ^initial_ash_bindings.current,
                on: true
              )

            {:left, true, false} ->
              from(_ in query,
                left_join: destination in ^relationship_destination,
                as: ^initial_ash_bindings.current,
                on: true
              )

            {:left, true, true} ->
              from(_ in query,
                left_lateral_join: destination in ^relationship_destination,
                as: ^initial_ash_bindings.current,
                on: true
              )

            {:left, false, false} ->
              from(_ in query,
                left_join: destination in ^relationship_destination,
                as: ^initial_ash_bindings.current,
                on:
                  field(as(^current_binding), ^relationship.source_attribute) ==
                    field(
                      destination,
                      ^relationship.destination_attribute
                    )
              )

            {:left, false, true} ->
              from(_ in query,
                left_lateral_join: destination in ^relationship_destination,
                as: ^initial_ash_bindings.current,
                on: true
              )
          end

        query
        |> AshSql.Aggregate.add_aggregates(
          used_aggregates,
          relationship.destination,
          false,
          initial_ash_bindings.current,
          {query.__ash_bindings__.resource, full_path}
        )
    end
  end

  @doc false
  def maybe_apply_filter(query, _root_query, _bindings, nil),
    do: {query, %AshSql.Expr.ExprInfo{}}

  def maybe_apply_filter(query, root_query, bindings, filter) do
    {dynamic, acc} = AshSql.Expr.dynamic_expr(root_query, filter, bindings, true)
    {from(row in query, where: ^dynamic), acc}
  end
end