/*
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements.  See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF licenses this file to You under the Apache License, Version 2.0
 * (the "License"); you may not use this file except in compliance with
 * the License.  You may obtain a copy of the License at
 *
 *    http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package org.apache.spark.sql.catalyst.analysis.resolver

import scala.collection.mutable

import org.apache.spark.SparkException
import org.apache.spark.sql.catalyst.analysis.UnresolvedHaving
import org.apache.spark.sql.catalyst.expressions.{Alias, Expression, NamedExpression}
import org.apache.spark.sql.catalyst.plans.logical.{Aggregate, Filter, LogicalPlan, Project}

/**
 * Resolves [[UnresolvedHaving]] node and its condition.
 */
class HavingResolver(resolver: Resolver, expressionResolver: ExpressionResolver)
    extends TreeNodeResolver[UnresolvedHaving, LogicalPlan]
    with RewritesAliasesInTopLcaProject
    with ResolvesNameByHiddenOutput
    with ValidatesFilter {
  private val scopes: NameScopeStack = resolver.getNameScopes
  private val expressionIdAssigner = expressionResolver.getExpressionIdAssigner
  private val autoGeneratedAliasProvider = new AutoGeneratedAliasProvider(expressionIdAssigner)

  /**
   * Resolve [[UnresolvedHaving]] operator:
   *  1. Resolve its child using the [[Resolver.resolve]].
   *  2. After resolving the child, transform the operator from [[UnresolvedHaving]] to [[Filter]].
   *     This is needed because result of the resolution is always a [[Filter]] (either on
   *     [[Aggregate]] or on [[Project]] depending on whether there are LCAs in the underlying
   *     [[Aggregate]]) and with this transformation we can avoid unnecessary checks for
   *     [[UnresolvedHaving]] in the rest of the code. Note that there is no `having` in Dataframes
   *     so unusual patterns are not expected.
   *  3. Resolve its condition using the [[ExpressionResolver.resolveExpressionTreeInOperator]].
   *     See `handleAggregateBelowHaving` doc for more details.
   *  4. Validate the filter using the [[ValidatesFilter.validateFilter]].
   */
  override def resolve(unresolvedHaving: UnresolvedHaving): LogicalPlan = {
    val resolvedChild = resolver.resolve(unresolvedHaving.child)

    val partiallyResolvedHaving =
      Filter(condition = unresolvedHaving.havingCondition, child = resolvedChild)
    val partiallyResolvedCondition = expressionResolver.resolveExpressionTreeInOperator(
      partiallyResolvedHaving.condition,
      partiallyResolvedHaving
    )

    val (resolvedCondition, missingExpressions) = resolvedChild match {
      case _ @(_: Project | _: Aggregate) if scopes.current.baseAggregate.isDefined =>
        handleAggregateBelowHaving(
          scopes.current.baseAggregate.get,
          partiallyResolvedCondition
        )
      case other =>
        throw SparkException.internalError(
          s"Unexpected operator ${other.getClass.getSimpleName} under HAVING"
        )
    }

    val (resolvedConditionWithAliasReplacement, filteredMissingExpressions) =
      tryReplaceSortOrderOrHavingConditionWithAlias(resolvedCondition, scopes, missingExpressions)

    val deduplicatedMissingExpressions = deduplicateMissingExpressions(filteredMissingExpressions)

    val resolvedChildWithMissingAttributes =
      insertMissingExpressions(resolvedChild, deduplicatedMissingExpressions)

    val isChildChangedByMissingExpressions = !resolvedChildWithMissingAttributes.eq(resolvedChild)

    val (finalChild, finalCondition) = resolvedChildWithMissingAttributes match {
      case project: Project if scopes.current.baseAggregate.isDefined =>
        val (newProject, newExpressions) = rewriteNamedExpressionsInTopLcaProject(
          projectToRewrite = project,
          baseAggregate = scopes.current.baseAggregate.get,
          expressionsToRewrite = Seq(resolvedConditionWithAliasReplacement),
          rewriteCandidates = missingExpressions,
          autoGeneratedAliasProvider = autoGeneratedAliasProvider
        )
        (newProject, newExpressions.head)
      case other => (other, resolvedCondition)
    }

    val resolvedHaving = partiallyResolvedHaving.copy(
      child = finalChild,
      condition = finalCondition
    )

    validateFilter(
      invalidExpressions = expressionResolver.getLastInvalidExpressionsInTheContextOfOperator,
      unresolvedOperator = unresolvedHaving,
      resolvedFilter = resolvedHaving
    )

    if (isChildChangedByMissingExpressions) {
      retainOriginalOutput(
        operator = resolvedHaving,
        missingExpressions = missingExpressions.toSeq,
        scopes = scopes
      )
    } else {
      resolvedHaving
    }
  }

  /**
   * Expression that are not resolved using the main output have to added to operators below
   * (and original output should be retained by placing a [[Project]] operator above the
   * [[Filter]] operator) for which we have three cases:
   *
   *  1. [[AggregateExpression]] not present in the [[Aggregate]]: added to the list while
   *     extracting the aggregate and grouping expressions in the
   *     [[GroupingAndAggregateExpressionsExtractor]].
   *     Example:
   *
   *     {{{ SELECT col1 FROM VALUES(1, 2) GROUP BY col1 HAVING max(col2) > 1; }}}
   *
   *     Plan would be:
   *     Project [col1#1]
   *     +- Filter (max(col2)#2 > 1)
   *        +- Aggregate [col1#1], [col1#1, max(col2#3) AS max(col2#3)#2]
   *           +- LocalRelation [col1#1, col2#3]
   *
   *  2. Expression is resolved using the [[Aggregate.groupingExpressions]]. These expression list
   *     is created while extracting grouping and aggregate expressions in the
   *     [[GroupingAndAggregateExpressionsExtractor]] by the [[UnresolvedHaving.havingCondition]].
   *     Example:
   *
   *     {{{ SELECT col1 FROM VALUES(1, 2) GROUP BY col1, col2 + 1 order by col2 + 1; }}}
   *
   *     Plan would be:
   *     Project [col1#1]
   *     +- Sort [(col2#2 + 1)#3 ASC NULLS FIRST], true
   *        +- Aggregate [col1#1, (col2#2 + 1)], [col1#1, (col2#2 + 1) AS (col2#2 + 1)#3]
   *           +- LocalRelation [col1#1, col2#2]
   *
   *  3. Outer references from the HAVING condition that are not present in the [[Aggregate]].
   *     Example:
   *
   *     {{{ SELECT col1 FROM VALUES(1) GROUP BY col1 HAVING 1 IN (SELECT MIN(col1)); }}}
   *
   *     Plan would be:
   *     Project [col1#1]
   *     +- Filter 1 IN (list#3 [min(col1#1)#2])
   *        :  +- Aggregate [outer(min(col1#1)#2) AS min(outer(col1))#4]
   *        :     +- OneRowRelation
   *     +- Aggregate [col1#1], [col1#1, min(col1#1) AS min(col1#1)#2]
   *        +- LocalRelation [col1#1]
   *
   * For more details on are missing expressions insert into the plan, see
   * [[ResolvesNameByHiddenOutput]] doc.
   */
  private def handleAggregateBelowHaving(
      aggregate: Aggregate,
      resolvedCondition: Expression): (Expression, Seq[NamedExpression]) = {
    val groupingAndAggregateExpressionsExtractor =
      new GroupingAndAggregateExpressionsExtractor(aggregate, autoGeneratedAliasProvider)

    val (finalCondition, extractedExpressions) = extractReferencedGroupingAndAggregateExpressions(
      resolvedCondition,
      groupingAndAggregateExpressionsExtractor
    )

    val lowerAliasedAggregateExpressions =
      resolver.getSubqueryRegistry.currentScope.getLowerOuterAggregateExpressionAliases.map {
        alias =>
          expressionIdAssigner.mapExpression(alias)
      }

    val lowerAliasedAggregateExpressionsFiltered = lowerAliasedAggregateExpressions.filterNot {
      alias =>
        groupingAndAggregateExpressionsExtractor.exists(alias.child)
    }

    (
      finalCondition,
      extractedExpressions ++
      lowerAliasedAggregateExpressionsFiltered
    )
  }

  /**
   * Method used for extraction of the grouping and aggregate expressions based on the HAVING
   * condition. This is done in a top-down manner by traversing the expression tree of the
   * condition, swapping an underlying expression found in the grouping or aggregate expressions
   * with the one that matches it and populating the `referencedGroupingExpressions` and
   * `extractedAggregateExpressionAliases` lists to insert missing expressions later.
   */
  private def extractReferencedGroupingAndAggregateExpressions(
      condition: Expression,
      groupingAndAggregateExpressionsExtractor: GroupingAndAggregateExpressionsExtractor)
      : (Expression, Seq[NamedExpression]) = {
    val referencedGroupingExpressions = new mutable.ArrayBuffer[NamedExpression]
    val extractedAggregateExpressionAliases = new mutable.ArrayBuffer[Alias]

    val resolvedCondition = condition.transformDown {
      case expression =>
        groupingAndAggregateExpressionsExtractor
          .extractReferencedGroupingAndAggregateExpressions(
            expression = expression,
            referencedGroupingExpressions = referencedGroupingExpressions,
            extractedAggregateExpressionAliases = extractedAggregateExpressionAliases
          )
    }

    (
      resolvedCondition,
      extractedAggregateExpressionAliases.toSeq ++ referencedGroupingExpressions.toSeq
    )
  }
}
