/*
 * 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, HashMap, LinkedHashSet}

import scala.collection.mutable.ArrayBuffer
import scala.jdk.CollectionConverters._

import org.apache.spark.SparkException
import org.apache.spark.sql.catalyst.analysis.AnalysisErrorAt
import org.apache.spark.sql.catalyst.expressions.{
  Alias,
  Expression,
  ExprId,
  ExprUtils,
  NamedExpression,
  ScalarSubquery
}
import org.apache.spark.sql.catalyst.expressions.aggregate.AggregateExpression
import org.apache.spark.sql.catalyst.plans.logical.{Aggregate, LogicalPlan, Project}
import org.apache.spark.sql.errors.QueryErrorsBase

/**
 * Handles resolution of lateral column references in [[Project]] and [[Aggregate]] operators.
 */
class LateralColumnAliasResolver(expressionResolver: ExpressionResolver) extends QueryErrorsBase {
  private val scopes = expressionResolver.getNameScopes
  private val aliasProvider = new AutoGeneratedAliasProvider(
    expressionResolver.getExpressionIdAssigner
  )

  /**
   * Handles lateral column references in [[Aggregate]] operator by doing the following:
   *  - Replaces the original [[Aggregate]] with a Project(Aggregate) structure in order to
   *  eliminate lateral column references from the [[Aggregate]] (detailed description can be found
   *  in `extractLcaAndReplaceAggWithProject`).
   *  - Asserts that the newly created [[Aggregate]] is valid.
   *  - Because the new [[Project]] now contains lateral column references, call
   *  `buildProjectWithResolvedLca` to resolve LCAs in the [[Project]].
   *  - Set the flag in the current [[NameScope]] that we have an LCA in [[Aggregate]] operator.
   *  - In order to be able to resolve [[Sort]] on top of an [[Aggregate]] that has LCAs, we need
   *  to collect all aliases from [[Aggregate]], as well as any aliases from artificially inserted
   *  [[Project]] nodes.
   *  - Collects all aliases from newly created [[Aggregate]] and [[Project]] nodes and adds them
   *  to `aliasesToCollect`
   */
  def handleLcaInAggregate(resolvedAggregate: Aggregate): AggregateWithLcaResolutionResult = {
    val aliasesToCollect = new ArrayBuffer[Alias]

    extractLcaAndReplaceAggWithProject(resolvedAggregate, aliasesToCollect) match {
      case _ @Project(projectList: Seq[_], aggregate: Aggregate) =>
        // TODO: This validation function does a post-traversal. This is discouraged in single-pass
        //       Analyzer.
        ExprUtils.assertValidAggregation(aggregate)

        val remappedAliases = new HashMap[ExprId, Alias](projectList.size)
        projectList.foreach {
          case alias: Alias => remappedAliases.put(alias.exprId, alias)
          case _ =>
        }

        val finalProject = buildProjectWithResolvedLca(
          resolvedChild = aggregate,
          scope = scopes.current,
          originalProjectList = projectList,
          firstIterationProjectList = aggregate.aggregateExpressions.map(_.toAttribute),
          remappedAliases = Some(remappedAliases),
          aliasesToCollect = aliasesToCollect
        )

        scopes.overwriteCurrent(output = Some(finalProject.projectList.map(_.toAttribute)))

        AggregateWithLcaResolutionResult(
          resolvedOperator = finalProject,
          outputList = finalProject.projectList,
          aggregateListAliases = aliasesToCollect.toSeq,
          baseAggregate = aggregate
        )
      case _ =>
        throw SparkException.internalError(
          s"Single-pass analyzer failed to resolve lateral column aliases in aggregate:" +
          s" ${resolvedAggregate.treeString}"
        )
    }
  }

  /**
   * Builds a multi-level [[Project]] with all lateral column aliases and their dependencies. First,
   * from top scope, we acquire dependency levels of all aliases. Dependency level is defined as a
   * number of attributes that an attribute depends on in the lateral alias reference chain. For
   * example, in a query like:
   *
   * {{{ SELECT 0 AS a, 1 AS b, 2 AS c, b AS d, a AS e, d AS f, a AS g, g AS h, h AS i }}}
   *
   * Dependency levels will be as follows:
   *
   * level 0: a, b, c
   * level 1: d, e, g
   * level 2: f, h
   * level 3: i
   *
   * Once we have dependency levels, we construct a multi-level [[Project]] in a following way:
   *  - There is exactly one [[Project]] node per level.
   *  - Project lists are compounded such that project lists on upper levels must contain all
   *  attributes from the below levels.
   *  - `iterationStart` is the project list on level 0.
   *  - In each iteration, if we are resolving LCAs from [[Aggregate]], we need to update aliases
   *  in the project list based on `remappedAliases` (detailed description in `updateAliasTrees`).
   *  `remappedAliases` should be None if we are resolving LCAs in [[Project]].
   *  - `originalProjectList` is placed on top of the multi-level [[Project]]. Any aliases that have
   *  been laterally referenced need to be replaced with only their names. This is because their
   *  full definitions ( `attr` as `name` ) have already been defined on lower levels.
   *  - If an attribute is never referenced, it does not show up in multi-level project lists, but
   *  instead only in the top-most [[Project]].
   *  - Additionally, collect all aliases from newly created [[Project]] nodes and add them to
   *  `aliasesToCollect`.
   *
   *  For previously given query, following above rules, resolved [[Project]] would look like:
   *
   * Project [a, b, 2 AS c, d, a AS e, d AS f, g, h, h AS i]
   * +- Project [a, b, d, g, g AS h]
   *    +- Project [a, b, b AS d, a AS g]
   *       +- Project [0 AS a, 1 AS b]
   *          +- OneRowRelation
   */
  def buildProjectWithResolvedLca(
      resolvedChild: LogicalPlan,
      scope: NameScope,
      originalProjectList: Seq[NamedExpression],
      firstIterationProjectList: Seq[NamedExpression],
      remappedAliases: Option[HashMap[ExprId, Alias]] = None,
      aliasesToCollect: ArrayBuffer[Alias] = ArrayBuffer.empty): Project = {
    val aliasDependencyMap = scope.lcaRegistry.getAliasDependencyLevels()
    val (finalChildPlan, _) = aliasDependencyMap.asScala.foldLeft(
      (resolvedChild, firstIterationProjectList)
    ) {
      case ((currentPlan, currentProjectList), availableAliases) =>
        val referencedAliases = new ArrayBuffer[Alias]
        val updatedAliases = updateAliasTrees(availableAliases, remappedAliases)
        updatedAliases.forEach(
          alias =>
            if (scope.lcaRegistry.isAttributeLaterallyReferenced(alias.toAttribute)) {
              referencedAliases.append(alias)
            }
        )

        if (referencedAliases.nonEmpty) {
          val newProjectList = currentProjectList.map(_.toAttribute) ++ referencedAliases

          newProjectList.foreach {
            case alias: Alias => aliasesToCollect += alias
            case _ =>
          }

          (Project(newProjectList, currentPlan), newProjectList)
        } else {
          (currentPlan, currentProjectList)
        }
    }

    val finalProjectList = originalProjectList.map { alias =>
      if (scope.lcaRegistry.isAttributeLaterallyReferenced(alias.toAttribute)) {
        alias.toAttribute
      } else {
        alias
      }
    }

    finalProjectList.foreach {
      case alias: Alias => aliasesToCollect += alias
      case _ =>
    }

    Project(finalProjectList, finalChildPlan)
  }

  /**
   * Replaces an [[Aggregate]] that has lateral column references with a [[Project]] by doing the
   * following:
   *  - Extracts the aggregate and grouping expressions from the [[Aggregate]] and places them in a
   *  new [[Aggregate]] operator with aliases on top of them. If an expression is already a
   *  [[NamedExpression]] we don't need to alias it again.
   *  - Places a [[Project]] on top of the new [[Aggregate]] operator, where the project list will
   *  be created from [[Alias]] references to original aggregate expressions.
   *  - Additionally, collect aliases from newly created aggregate expressions and add them to
   *  `aliasesToCollect`.
   *
   * For example, for a query like:
   *
   * {{{ SELECT 1 AS a, a + 1 AS b, col1, MAX(col1) FROM VALUES(1) GROUP BY col1 }}}
   *
   * the result of extraction will be:
   *
   * !Project [1 AS a#4, (a#4 + 1) AS b#5, col1#3, max(col1)#11 AS max(col1)#10]
   *  +- Aggregate [col1#3], [col1#3, max(col1#3) AS max(col1)#11]
   *     +- LocalRelation [col1#3]
   *
   * The [[Project]] is unresolved, which is fine, because it will later be resolved as if we only
   * had a lateral alias reference in [[Project]] and not [[Aggregate]].
   */
  private def extractLcaAndReplaceAggWithProject(
      aggregate: Aggregate,
      aliasesToCollect: ArrayBuffer[Alias]): Project = {
    val newAggregateExpressions = new LinkedHashSet[NamedExpression]
    val extractedExpressionAliases = new HashMap[Expression, NamedExpression]()
    val groupingExpressionSemanticComparator = new SemanticComparator(aggregate.groupingExpressions)
    val extractedExpressions = aggregate.aggregateExpressions
      .map(
        aggregateExpression =>
          extractAggregateAndGroupingExpressions(
            expression = aggregateExpression,
            groupingExpressionsSemanticComparator = groupingExpressionSemanticComparator,
            extractedExpressionAliases = extractedExpressionAliases,
            newAggregateExpressions = newAggregateExpressions
          ).asInstanceOf[NamedExpression]
      )

    val newAggregateExpressionsSeq = newAggregateExpressions.asScala.toSeq
    newAggregateExpressionsSeq.foreach {
      case alias: Alias => aliasesToCollect += alias
      case _ =>
    }

    val result = Project(
      projectList = extractedExpressions,
      child = aggregate.copy(aggregateExpressions = newAggregateExpressionsSeq)
    )
    result
  }

  /**
   * Extracts aggregate and grouping expressions from an [[Aggregate]] that has lateral column
   * references in order to place them the newly created [[Aggregate]] without LCAs.
   */
  private def extractAggregateAndGroupingExpressions(
      expression: Expression,
      groupingExpressionsSemanticComparator: SemanticComparator,
      extractedExpressionAliases: HashMap[Expression, NamedExpression],
      newAggregateExpressions: LinkedHashSet[NamedExpression]): Expression = {
    expression match {
      case aggregateExpression: AggregateExpression =>
        val ne = getOrUpdateMapping(extractedExpressionAliases, aggregateExpression)
        newAggregateExpressions.add(ne)
        ne.toAttribute
      case expression if groupingExpressionsSemanticComparator.exists(expression) =>
        val ne = getOrUpdateMapping(extractedExpressionAliases, expression)
        newAggregateExpressions.add(ne)
        ne.toAttribute
      case scalarSubquery: ScalarSubquery
          if scalarSubquery.children.nonEmpty && !groupingExpressionsSemanticComparator.exists(
            scalarSubquery
          ) =>
        scalarSubquery.failAnalysis(
          errorClass = "SCALAR_SUBQUERY_IS_IN_GROUP_BY_OR_AGGREGATE_FUNCTION",
          messageParameters = Map("sqlExpr" -> toSQLExpr(scalarSubquery))
        )
      case expression =>
        expression.mapChildren(
          child =>
            extractAggregateAndGroupingExpressions(
              expression = child,
              groupingExpressionsSemanticComparator = groupingExpressionsSemanticComparator,
              extractedExpressionAliases = extractedExpressionAliases,
              newAggregateExpressions = newAggregateExpressions
            )
        )
    }
  }

  /**
   * Gets the [[Alias]] mapping for the given expression or assigns a new [[Alias]] through
   * [[AutoGeneratedAliasProvider]] if it doesn't exist.
   */
  private def getOrUpdateMapping(
      extractedExpressionAliases: HashMap[Expression, NamedExpression],
      expression: Expression): NamedExpression = {
    val key = expression.canonicalized
    extractedExpressionAliases.computeIfAbsent(
      key,
      (_: Expression) =>
        expression match {
          case namedExpression: NamedExpression => namedExpression
          case _ => aliasProvider.newAlias(expression)
        }
    )
  }

  /**
   * Updates expression trees below aliases based on `remappedAliases`. This is necessary in order
   * to stay compatible with fixed-point when resolving LCAs in [[Aggregate]].
   * For example, consider the following query:
   *
   * {{{ SELECT 1 AS a, a + 1 AS b, MAX(col1) FROM VALUES(1) GROUP BY col1 }}}
   *
   * The unresolved plan will be:
   *
   * !Aggregate [col1#3], [1 AS a#4, (a#4 + 1) AS b#5, max(col1#3) AS max(col1)#10]
   * +- LocalRelation [col1#3]
   *
   * and after resolving all LCAs we will have:
   *
   * Project [a#4, (a#4 + 1) AS b#5, max(col1)#11 AS max(col1)#10]
   * +- Project [max(col1)#11, 1 AS a#4]
   *    +- Aggregate [col1#3], [max(col1#3) AS max(col1)#11]
   *       +- LocalRelation [col1#3]
   *
   * An unresolved [[Aggregate]] will contain an aggregate expression `max(col1#3) as max(col1)#10`
   * with its implicit [[Alias]]. In order to preserve the [[Alias]] with [[ExprId]] #2 we need
   * to place a new [[Alias]] (`max(col1)#11`) on top of the aggregate expression (this is done in
   * `extractLcaAndReplaceAggWithProject`). Now we can place the resulting [[Alias]] in the newly
   * created [[Aggregate]] that contains only aggregate and grouping expressions. However, in order
   * to reference this new [[Alias]] from the [[Project]] that we place on top of the new
   * [[Aggregate]], we need to update the expression tree of `max(col1)#10` [[Alias]] to point to
   * the newly created alias `max(col1)#11` and we get `max(col1)#11 as max(col1)#10` in the final
   * project list.
   */
  private def updateAliasTrees(
      aliases: ArrayList[Alias],
      remappedAliases: Option[HashMap[ExprId, Alias]]): ArrayList[Alias] = {
    remappedAliases match {
      case Some(map) =>
        val updatedAliases = new ArrayList[Alias](aliases.size())
        aliases.forEach(alias => {
          val updatedAlias = map.getOrDefault(alias.exprId, alias)
          updatedAliases.add(updatedAlias)
        })
        updatedAliases
      case None => aliases
    }
  }
}
