/*
 * 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 org.apache.spark.sql.AnalysisException
import org.apache.spark.sql.catalyst.analysis.TypeCoercionValidation
import org.apache.spark.sql.catalyst.expressions.{
  Expression,
  OuterReference,
  SubExprUtils
}
import org.apache.spark.sql.catalyst.expressions.aggregate.{
  AggregateExpression,
  AggregateFunction,
  ListAgg
}
import org.apache.spark.sql.catalyst.util.toPrettySQL
import org.apache.spark.sql.errors.QueryCompilationErrors

/**
 * Resolver for [[AggregateExpressions]] that can come from either [[FunctionResolver]] or
 * [[ExpressionResolver]]. It handles the resolution and validation of [[AggregateExpression]].
 */
class AggregateExpressionResolver(
    operatorResolver: Resolver,
    expressionResolver: ExpressionResolver)
    extends TreeNodeResolver[AggregateExpression, Expression]
    with ResolvesExpressionChildren
    with CoercesExpressionTypes {

  private val traversals = expressionResolver.getExpressionTreeTraversals

  private val expressionResolutionContextStack =
    expressionResolver.getExpressionResolutionContextStack
  private val subqueryRegistry = operatorResolver.getSubqueryRegistry
  private val autoGeneratedAliasProvider = expressionResolver.getAutoGeneratedAliasProvider

  /**
   * Resolves the given [[AggregateExpression]] originating from [[ExpressionResolver]] by
   * resolving its children recursively and validating the resolved expression.
   */
  override def resolve(aggregateExpression: AggregateExpression): Expression = {
    val expressionResolutionContext = expressionResolutionContextStack.peek()
    expressionResolutionContext.resolvingTreeUnderAggregateExpression = true
    if (expressionResolutionContext.resolvingWindowFunction) {
      expressionResolutionContext.windowFunctionNestednessLevel += 1
    }
    val aggregateExpressionWithChildrenResolved =
      withResolvedChildren(aggregateExpression, expressionResolver.resolve _)
        .asInstanceOf[AggregateExpression]
    handleAggregateExpressionWithChildrenResolved(aggregateExpressionWithChildrenResolved)
  }

  /**
   * Resolves the given [[AggregateExpression]] originating from [[FunctionResolver]] by applying
   * type coercion to its children and validating the resolved expression. In this case, it is not
   * necessary to resolve the children recursively, as they were already resolved in
   * [[FunctionResolver]].
   */
  def resolveWithoutRecursingIntoChildren(aggregateExpression: AggregateExpression): Expression = {
    val aggregateExpressionWithTypeCoercedChildren = aggregateExpression
      .mapChildren(
        expression => coerceExpressionTypes(expression, traversals.current)
      )
      .asInstanceOf[AggregateExpression]
    handleAggregateExpressionWithChildrenResolved(aggregateExpressionWithTypeCoercedChildren)
  }

  /**
   * Handles resolution and validation of the [[AggregateExpression]] after its children have been
   * resolved:
   *  - Resolution:
   *    1. Update the [[ExpressionResolver.expressionResolutionContextStack]];
   *    2. Handle [[OuterReference]] in [[AggregateExpression]], if there are any (see
   *       `handleOuterAggregateExpression`);
   *    3. Confirm that we are resolving an actual aggregate expression, not a window aggregate
   *       function by assessing the `windowFunctionNestednessLevel`:
   *
   *          - In this query, `windowFunctionNestednessLevel` is `1` when resolving `SUM()`,
   *            meaning we are resolving a window aggregate function:
   *
   *            {{{ SELECT SUM() OVER() FROM VALUES 1, 2; }}}
   *
   *          - In this query, `windowFunctionNestednessLevel` is `2` when resolving `SUM(col1)`,
   *            meaning we are resolving an aggregate expression:
   *
   *            {{{
   *            SELECT
   *              SUM(SUM(col1)) OVER (PARTITION BY SUM(col1))
   *            FROM
   *              VALUES (1, 2)
   *            GROUP BY col2;
   *            }}}
   *
   *          - In this query, `windowFunctionNestednessLevel` is `0` when resolving `SUM(col1)`,
   *            meaning we are resolving an aggregate expression:
   *
   *            {{{ SELECT SUM(col1) FROM VALUES (1, 2); }}}
   *
   *       If the expression is not a window aggregate function, set
   *       `hasAggregateExpressionsOutsideWindow` to `true` and collect it as window source
   *       expression.
   *
   *       In this query:
   *
   *       {{{ SELECT SUM(SUM(col1)) OVER () FROM VALUES (1, 2) GROUP BY col2; }}}
   *
   *       `SUM(col1)` is collected, whereas `SUM(SUM(col1))` is not because the latter is a
   *       window aggregate function.  As a result, `SUM(col1)` is handled by [[Aggregate]],
   *       while `SUM(SUM(col1))` is handled  by [[Window]]:
   *
   *       Window [sum(_w0#2) windowspecdefinition...#3]
   *       +- Aggregate [col2#1], [sum(col1#0) AS _w0#2]
   *  - Validation:
   *   1. Validate that the type coercion is doen properly for the underlying aggregate function.
   *   2. [[ListAgg]] is not allowed in DISTINCT aggregates if the order value is ambiguous
   *      after deduplication. When [[SQLConf.LISTAGG_ALLOW_DISTINCT_CAST_WITH_ORDER]] is
   *      enabled, a mismatch is tolerated if the child value uniquely determines the order
   *      value (see [[ListAgg.hasDistinctOrderAmbiguity]]);
   *   3. [[PythonUDAF]] is not allowed in [[Pivot.aggregates]];
   *   4. Nested aggregate functions are not allowed, unless the outer function represents
   *      [[WindowExpression.windowFunction]]. For example the following query is considered
   *      valid, because `SUM(SUM(col1))` is followed by an `OVER` clause, making it a nesting
   *      of aggregate expression in a window expression:
   *
   *      {{{
   *      SELECT
   *      SUM(SUM(col1)) OVER (PARTITION BY col2)
   *      FROM (VALUES (1, 'a'), (2, 'b'))
   *      GROUP BY col2
   *      }}}
   *
   *      On the other hand, this query should be invalid, because `SUM(SUM(col1))` isn't followed
   *      by an `OVER` clause, therefore it should be treated like aggregate expression nesting:
   *
   *      {{{
   *      SELECT
   *      SUM(SUM(SUM(col1)) OVER (PARTITION BY col2)
   *      FROM (VALUES (1, 'a'), (2, 'b'))
   *      GROUP BY col2
   *      }}}
   *
   *   5. Nondeterministic expressions in the subtree of a related aggregate function are not
   *      allowed;
   *   6. The mix of outer and local references is not allowed;
   */
  private def handleAggregateExpressionWithChildrenResolved(
      aggregateExpressionWithChildrenResolved: AggregateExpression): Expression = {
    val expressionResolutionContext = expressionResolutionContextStack.peek()

    validateResolvedAggregateExpression(aggregateExpressionWithChildrenResolved)

    // There are two different cases that we handle regarding the value of the flag:
    //
    //   - We have an attribute under an `AggregateExpression`:
    //       {{{ SELECT COUNT(col1) FROM VALUES (1); }}}
    //     In this case, value of the `hasAttributeOutsideOfAggregateExpressions` flag should be
    //     `false` as it indicates whether there is an attribute in the subtree that's not
    //     `AggregateExpression` so we can throw the `MISSING_GROUP_BY` exception appropriately.
    //
    //   - In the following example:
    //       {{{ SELECT COUNT(*), col1 + 1 FROM VALUES (1); }}}
    //     It would be `true` as described above.
    expressionResolutionContext.hasAttributeOutsideOfAggregateExpressions = false

    val resolvedAggregateExpression = if (expressionResolutionContext.hasOuterReferences) {
      handleOuterAggregateExpression(aggregateExpressionWithChildrenResolved)
    } else {
      aggregateExpressionWithChildrenResolved
    }

    if (expressionResolutionContext.windowFunctionNestednessLevel != 1) {
      expressionResolutionContext.hasAggregateExpressionsOutsideWindow = true
    }

    resolvedAggregateExpression
  }

  private def validateResolvedAggregateExpression(aggregateExpression: AggregateExpression): Unit =
    aggregateExpression match {
      case _ @AggregateExpression(aggregateFunction: AggregateFunction, _, _, _, _)
          if aggregateFunction.checkInputDataTypes().isFailure =>
        TypeCoercionValidation.failOnTypeCheckResult(aggregateFunction)
      case agg @ AggregateExpression(listAgg: ListAgg, _, _, _, _)
          if agg.isDistinct && listAgg.hasDistinctOrderAmbiguity =>
        listAgg.throwDistinctOrderError()
      case _ =>
        val expressionResolutionContext = expressionResolutionContextStack.peek()
        if (expressionResolutionContext.windowFunctionNestednessLevel != 1 &&
          expressionResolutionContext.hasAggregateExpressionsOutsideWindow) {
          throwNestedAggregateFunction(aggregateExpression)
        }

        aggregateExpression.aggregateFunction.children.foreach { child =>
          if (!child.deterministic) {
            throwAggregateFunctionWithNondeterministicExpression(
              aggregateExpression,
              child
            )
          }
        }
    }

  /**
   * If the [[AggregateExpression]] has outer references in its subtree, we need to handle it in a
   * special way. The whole process is explained in the [[SubqueryScope]] scaladoc, but in short
   * we need to:
   *  - Validate that we don't have local references in this subtree;
   *  - Create a new subtree without [[OuterReference]]s;
   *  - Alias this subtree and put it inside the current [[SubqueryScope]];
   *  - If outer aggregates are allowed, replace the [[AggregateExpression]] with an
   *    [[OuterReference]] to the auto-generated [[Alias]] that we created in case the subtree
   *    without [[OuterReference]]s can't be found in the outer
   *    [[Aggregate.aggregateExpressions]] list or a semantically equal expression wasn't already
   *    added. Otherwise, use the [[Alias]] from the outer [[Aggregate]]. This alias will later be
   *    injected into the outer [[Aggregate]];
   *  - Store the name that needs to be used for the [[OuterReference]] in
   *    [[OuterReference.SINGLE_PASS_SQL_STRING_OVERRIDE]] and
   *    [[OuterReference.SINGLE_PASS_OUTER_AGGREGATE_ALIAS_NAME_OVERRIDE]], both computed from the
   *    pretty-printed original [[AggregateExpression]] (before [[OuterReference]] stripping). The
   *    former overrides [[OuterReference.sql]] and the latter overrides the alias name assigned by
   *    [[AliasResolution.resolve]].
   *  - In case we have an [[AggregateExpression]] inside a [[Sort]] operator, we need to handle it
   *    in a special way (see [[handleAggregateExpressionOutsideAggregate]] for more details).
   *  - Return the original [[AggregateExpression]] otherwise. This is done to stay compatible
   *    with the fixed-point Analyzer - a proper exception will be thrown later by
   *    [[ValidateSubqueryExpression]].
   */
  private def handleOuterAggregateExpression(
      aggregateExpression: AggregateExpression): Expression = {
    if (expressionResolutionContextStack.peek().hasLocalReferences) {
      throw QueryCompilationErrors.mixedRefsInAggFunc(
        aggregateExpression.sql,
        aggregateExpression.origin
      )
    }

    val resolvedOuterAggregateExpression =
      if (subqueryRegistry.currentScope.aggregateExpressionsExtractor.isDefined) {
        extractOuterAggregateExpression(
          aggregateExpression = aggregateExpression,
          aggregateExpressionsExtractor =
            subqueryRegistry.currentScope.aggregateExpressionsExtractor.get
        )
      } else {
        aggregateExpression
      }

    resolvedOuterAggregateExpression match {
      case outerReference: OuterReference =>
        val name = toPrettySQL(aggregateExpression)
        outerReference.setTagValue(OuterReference.SINGLE_PASS_SQL_STRING_OVERRIDE, name)
        outerReference.setTagValue(
          OuterReference.SINGLE_PASS_OUTER_AGGREGATE_ALIAS_NAME_OVERRIDE,
          name
        )
        outerReference
      case other => other
    }
  }

  private def extractOuterAggregateExpression(
      aggregateExpression: AggregateExpression,
      aggregateExpressionsExtractor: GroupingAndAggregateExpressionsExtractor): OuterReference = {
    val aggregateExpressionWithStrippedOuterReferences =
      SubExprUtils.stripOuterReference(aggregateExpression)

    val outerAggregateExpressionAlias = autoGeneratedAliasProvider.newOuterAlias(
      child = aggregateExpressionWithStrippedOuterReferences
    )

    subqueryRegistry.currentScope.getExtractedAggregateExpression(
      aggregateExpressionWithStrippedOuterReferences
    ) match {
      case Some(alreadyExtractedAggregateExpression) =>
        OuterReference(alreadyExtractedAggregateExpression.toAttribute)
      case None =>
        val (_, referencedAggregateExpressionAlias) =
          aggregateExpressionsExtractor.collectFirstAggregateExpression(
            aggregateExpressionWithStrippedOuterReferences
          )

        referencedAggregateExpressionAlias match {
          case Some(alias) =>
            subqueryRegistry.currentScope.addExtractedAggregateExpression(alias)
            OuterReference(alias.toAttribute)
          case None =>
            subqueryRegistry.currentScope.addExtractedAggregateExpression(
              outerAggregateExpressionAlias
            )
            OuterReference(outerAggregateExpressionAlias.toAttribute)
        }
    }
  }

  private def throwNestedAggregateFunction(aggregateExpression: AggregateExpression): Nothing = {
    throw new AnalysisException(
      errorClass = "NESTED_AGGREGATE_FUNCTION",
      messageParameters = Map.empty,
      origin = aggregateExpression.origin
    )
  }

  private def throwAggregateFunctionWithNondeterministicExpression(
      aggregateExpression: AggregateExpression,
      nonDeterministicChild: Expression): Nothing = {
    throw new AnalysisException(
      errorClass = "AGGREGATE_FUNCTION_WITH_NONDETERMINISTIC_EXPRESSION",
      messageParameters = Map("sqlExpr" -> toSQLExpr(aggregateExpression)),
      origin = nonDeterministicChild.origin
    )
  }

}
