/*
 * 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 java.util.{ArrayList, LinkedHashMap}

import scala.jdk.CollectionConverters._

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, Window}

/**
 * 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 = expressionResolver.getAutoGeneratedAliasProvider
  private val operatorResolutionContextStack = resolver.getOperatorResolutionContextStack

  /**
   * Resolve [[UnresolvedHaving]] operator:
   *  1. Resolve its child using the [[Resolver.resolve]].
   *  2. Delegate resolution to appropriate method depending on `resolvedChild`'s type:
   *    - For [[Window]] with base [[Aggregate]], delegate resolution to `resolveHavingAboveWindow`.
   *    - Delegate resolution to `handleHavingWithResolvedChild` in all other cases.
   */
  override def resolve(unresolvedHaving: UnresolvedHaving): LogicalPlan = {
    val resolvedChild = resolver.resolve(unresolvedHaving.child)

    resolvedChild match {
      case window: Window if scopes.current.baseAggregate.isDefined =>
        resolveHavingAboveWindow(
          window = window,
          unresolvedHaving = unresolvedHaving,
          resolvedChild = resolvedChild
        )
      case _ =>
        handleHavingWithResolvedChild(
          unresolvedHaving = unresolvedHaving,
          resolvedChild = resolvedChild
        )
    }
  }

  /**
   * Find [[Window]]'s top-most [[Aggregate]] child and resolve [[UnresolvedHaving]] as if it was
   * placed directly on top of it in the parsed plan.
   *
   * Consider this query:
   *
   * {{{
   * SELECT
   *   col1,
   *   SUM(col1) OVER ()
   * FROM VALUES (1, 2)
   * GROUP BY col1
   * HAVING col1 > 0;
   * }}}
   *
   * Parsed plan:
   *
   * 'UnresolvedHaving ('col1 > 0)
   * +- 'Aggregate ['col1], ['col1, unresolvedalias('SUM('col1) windowspec...)]
   *
   * After resolving [[UnresolvedHaving]]'s child [[Aggregate]] node, it gets replaced by
   * Window -> Aggregate structure:
   *
   * 'UnresolvedHaving ('col1 > 0)
   * Window [col1#0, sum(col1#0) windowspec...#2]
   * +- Aggregate [col1#0], [col1#0]
   *
   * Since [[Filter]] is not dependant on the [[Window]], we push a [[Filter]] operator down to
   * the nearest [[Aggregate]] and perform the rest of resolution using its output:
   *
   * Window [col1#0, sum(col1#0) windowspec...#1]
   * +- Filter (col1#0 > 0)
   *    +- Aggregate [col1#0], [col1#0]
   *
   * Here's another query demonstrating multi-level [[Window]] nesting:
   *
   * {{{
   * SELECT
   *   col1,
   *   SUM(col1) OVER (),
   *   SUM(col1) OVER (ORDER BY col1)
   * FROM VALUES (1, 2)
   * GROUP BY col1
   * HAVING col1 > 0;
   * }}}
   *
   * Parsed plan:
   *
   * 'UnresolvedHaving ('col1 > 0)
   * +- 'Aggregate ['col1],
   *               ['col1,
   *               unresolvedalias('SUM('col1) windowspec...,
   *               unresolvedalias('SUM('col1) windowspec...)]
   *
   * Plan after resolving [[Aggregate]]:
   *
   * Window [col1#0, sum(col1#0) windowspec...#2, ...#1]
   * +- Window [col1#0, sum(col1#0) windowspec...#1], [col1#0 ASC NULLS FIRST]
   *    +- Aggregate [col1#0], [col1#0]
   *
   * Analyzed plan:
   *
   * Window [col1#0, sum(col1#0) windowspec...#2, ...#1]
   * +- Window [col1#0, sum(col1#0) windowspec...#1], [col1#0 ASC NULLS FIRST]
   *    +- Filter (col1#0 > 0)
   *       +- Aggregate [col1#0], [col1#0]
   *
   * 1. Preserve current output in `windowOutput`, because [[UnresolvedHaving]]'s parent should
   *    consume it. It will be altered while inserting [[Filter]] below [[Window]], since HAVING
   *    resolution is based on [[Aggregate]] output.
   * 2. Find the first [[Aggregate]] below a [[Window]] chain. Maintain a list of encountered
   *    [[Window]] operators.
   * 3. Place a [[Filter]] on top of [[Aggregate]] by resolving [[UnresolvedHaving]] as if it was
   *    placed directly above the [[Aggregate]]:
   *     - Overwrite current output with [[Aggregate]] output.
   *     - Resolve [[Filter]] using `handleHavingWithResolvedChild`.
   * 4. Reset output to `windowOutput` and place the original [[Window]] operators on top of the
   *    created [[Filter]].
   */
  private def resolveHavingAboveWindow(
      window: Window,
      unresolvedHaving: UnresolvedHaving,
      resolvedChild: LogicalPlan): Window = {
    val windowOutput = scopes.current.output

    var currentPlan: LogicalPlan = window
    val windowChain = new ArrayList[Window]()

    while (currentPlan.isInstanceOf[Window]) {
      val window = currentPlan.asInstanceOf[Window]
      windowChain.add(window)
      currentPlan = window.child
    }

    val filterWithAggregate = currentPlan match {
      case aggregate: Aggregate =>
        scopes.overwriteCurrent(
          output = Some(
            aggregate.aggregateExpressions.map(namedExpression => namedExpression.toAttribute)
          ),
          hiddenOutput = Some(scopes.current.hiddenOutput),
          availableAliases = Some(scopes.current.availableAliases),
          aggregateListAliases = scopes.current.aggregateListAliases,
          baseAggregate = Some(aggregate)
        )
        handleHavingWithResolvedChild(
          unresolvedHaving = unresolvedHaving,
          resolvedChild = aggregate
        )
      case other =>
        throw SparkException.internalError(
          s"Expected an Aggregate below a Window chain, but found ${other.getClass.getSimpleName}"
        )
    }

    scopes.overwriteCurrent(
      output = Some(windowOutput),
      hiddenOutput = Some(scopes.current.hiddenOutput),
      availableAliases = Some(scopes.current.availableAliases),
      aggregateListAliases = scopes.current.aggregateListAliases,
      baseAggregate = scopes.current.baseAggregate
    )

    windowChain.asScala
      .foldRight(filterWithAggregate) { (window, childPlan) =>
        window.withNewChildren(Seq(childPlan))
      }
      .asInstanceOf[Window]
  }

  /**
   * Helper method that resolves [[UnresolvedHaving]] operator whose child is already resolved.
   *
   *  1. 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.
   *  2. Resolve its condition using the [[ExpressionResolver.resolveExpressionTreeInOperator]].
   *     See `handleAggregateBelowHaving` doc for more details.
   *  3. Validate the filter using the [[ValidatesFilter.validateFilter]].
   */
  private def handleHavingWithResolvedChild(
      unresolvedHaving: UnresolvedHaving,
      resolvedChild: LogicalPlan): LogicalPlan = {
    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 resolvedChildWithMissingAttributes =
      insertMissingExpressions(resolvedChild, filteredMissingExpressions)

    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,
        operatorResolutionContextStack = operatorResolutionContextStack
      )
    } else {
      resolvedHaving
    }
  }

  /**
   * Expression that are not resolved using the main output have to be 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 aggregateExpressionsExtractedFromSubquery = new LinkedHashMap[Expression, Alias]()
    resolver.getOperatorResolutionContextStack.current.getAllSubqueryAggregateExpressions.foreach {
      case alias: Alias if !groupingAndAggregateExpressionsExtractor.exists(alias.child) =>
        val mappedAlias = expressionIdAssigner.mapExpression(alias)
        aggregateExpressionsExtractedFromSubquery.put(mappedAlias.child.canonicalized, mappedAlias)
        mappedAlias
      case _ =>
    }

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

    (
      finalCondition,
      extractedExpressions ++
      aggregateExpressionsExtractedFromSubquery.values().asScala.toSeq
    )
  }

  /**
   * Method used for extraction of the grouping and aggregate expressions, that have not already
   * been extracted from [[SubqueryExpression]], 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,
      aggregateExpressionsExtractedFromSubquery: LinkedHashMap[Expression, Alias])
      : (Expression, Seq[NamedExpression]) = {
    val referencedGroupingExpressions = new LinkedHashMap[Expression, NamedExpression]
    val extractedAggregateExpressionAliases = new LinkedHashMap[Expression, Alias]

    val resolvedCondition = condition.transformDown {
      case expression
          if aggregateExpressionsExtractedFromSubquery.containsKey(expression.canonicalized) =>
        aggregateExpressionsExtractedFromSubquery.get(expression.canonicalized).toAttribute
      case expression =>
        groupingAndAggregateExpressionsExtractor
          .extractReferencedGroupingAndAggregateExpressions(
            expression = expression,
            referencedGroupingExpressions = referencedGroupingExpressions,
            extractedAggregateExpressionAliases = extractedAggregateExpressionAliases
          )
    }

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