/*
 * 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.catalyst.analysis.{AliasResolution, UnresolvedAlias}
import org.apache.spark.sql.catalyst.expressions.{
  Alias,
  AliasHelper,
  Expression,
  NamedExpression,
  OuterReference
}
import org.apache.spark.sql.errors.QueryCompilationErrors

/**
 * Resolver class that resolves unresolved aliases and handles user-specified aliases.
 */
class AliasResolver(expressionResolver: ExpressionResolver)
    extends TreeNodeResolver[UnresolvedAlias, Expression]
    with ResolvesExpressionChildren
    with AliasHelper {
  private val scopes = expressionResolver.getNameScopes
  private val expressionResolutionContextStack =
    expressionResolver.getExpressionResolutionContextStack
  private val autoGeneratedAliasProvider = new AutoGeneratedAliasProvider(
    expressionResolver.getExpressionIdAssigner
  )

  /**
   * Resolves [[UnresolvedAlias]] by resolving its child and computing the alias name by calling
   * [[AliasResolution]] on the result. After resolving it, we assign a correct exprId to the
   * resulting [[Alias]]. In case result of the [[AliasResolution]] call is an [[OuterReference]],
   * we create a new [[Alias]] using the [[AutoGeneratedAliasProvider]]. Here we allow inner
   * aliases to persist until the end of single-pass resolution, after which they will be removed
   * in the post-processing phase.
   *
   * Resulting [[Alias]] must be added to the list of `availableAliases` in the current
   * [[NameScope]].
   */
  override def resolve(unresolvedAlias: UnresolvedAlias): NamedExpression =
    scopes.current.lcaRegistry.withNewLcaScope(
      isTopLevelAlias = expressionResolutionContextStack.peek().isTopOfProjectList
    ) {
      val aliasWithResolvedChildren =
        withResolvedChildren(unresolvedAlias, expressionResolver.resolve _)
          .asInstanceOf[UnresolvedAlias]

      val resolvedNode =
        AliasResolution.resolve(aliasWithResolvedChildren).asInstanceOf[NamedExpression]

      resolvedNode match {
        case alias: Alias =>
          val resultAlias = expressionResolver.getExpressionIdAssigner.mapExpression(alias)
          scopes.current.availableAliases.add(resultAlias.exprId)
          resultAlias
        case outerReference: OuterReference =>
          autoGeneratedAliasProvider.newAlias(outerReference)
        case _ =>
          throw QueryCompilationErrors.unsupportedSinglePassAnalyzerFeature(
            s"${resolvedNode.getClass} expression resolution"
          )
      }
    }

  /**
   * Handle already resolved [[Alias]] nodes, i.e. user-specified aliases. Here we only need to
   * resolve its children and afterwards reassign exprId to the resulting [[Alias]]. Resulting
   * [[Alias]] must be added to the list of `availableAliases` in the current [[NameScope]].
   *
   * When resolving grouping expressions, we may meet duplicate aliases from grouping expressions.
   * This is only the case with partially resolved DataFrame logical plans. We need to deprioritize
   * those aliases. See [[ExpressionIdAssigner.mapExpression]] doc for more details.
   */
  def handleResolvedAlias(alias: Alias): Alias = {
    val resolvedAlias = scopes.current.lcaRegistry.withNewLcaScope(
      isTopLevelAlias = expressionResolutionContextStack.peek().isTopOfProjectList
    ) {
      val aliasWithResolvedChildren =
        withResolvedChildren(alias, expressionResolver.resolve _).asInstanceOf[Alias]

      val resultAlias = expressionResolver.getExpressionIdAssigner.mapExpression(
        originalExpression = aliasWithResolvedChildren,
        prioritizeOldDuplicateAliasId =
          expressionResolutionContextStack.peek().resolvingGroupingExpressions
      )

      scopes.current.availableAliases.add(resultAlias.exprId)

      resultAlias
    }

    collapseAlias(resolvedAlias)
  }

  /**
   * In case where there are two explicit [[Alias]]es, one on top of the other, remove the bottom
   * one. For the example bellow:
   *
   * - df.select($"column".as("alias_1").as("alias_2"))
   *
   * the plan is:
   *
   *   Project[
   *     Alias("alias_2")(
   *       Alias("alias_1")(id1)
   *     )(id2)
   *   ]( ... )
   *
   * and after the `collapseAlias` call (removing the bottom one) it would be:
   *
   *   Project[
   *     Alias("alias_2")(id2)
   *   ]( ... )
   */
  private def collapseAlias(alias: Alias): Alias =
    alias.child match {
      case _: Alias => mergeAliases(alias)
      case _ => alias
    }
}
