/*
 * 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.UnresolvedInlineTable
import org.apache.spark.sql.catalyst.expressions.{
  Exists,
  Expression,
  InSubquery,
  ListQuery,
  ScalarSubquery,
  SubExprUtils
}
import org.apache.spark.sql.catalyst.plans.logical.LogicalPlan
import org.apache.spark.sql.errors.QueryCompilationErrors

/**
 * [[SubqueryExpressionResolver]] resolves specific [[SubqueryExpression]]s, such as
 * [[ScalarSubquery]], [[ListQuery]], and [[Exists]].
 */
class SubqueryExpressionResolver(expressionResolver: ExpressionResolver, resolver: Resolver)
    extends CoercesExpressionTypes {
  private val scopes = resolver.getNameScopes
  private val operatorResolutionContextStack = resolver.getOperatorResolutionContextStack
  private val traversals = expressionResolver.getExpressionTreeTraversals
  private val cteRegistry = resolver.getCteRegistry
  private val subqueryRegistry = resolver.getSubqueryRegistry
  private val expressionIdAssigner = expressionResolver.getExpressionIdAssigner
  private val autoGeneratedAliasProvider = expressionResolver.getAutoGeneratedAliasProvider
  private val expressionResolutionContextStack =
    expressionResolver.getExpressionResolutionContextStack

  /**
   * Resolve [[ScalarSubquery]]:
   *  - Resolve the subquery plan;
   *  - Get outer references;
   *  - Type coerce it;
   *  - Add it to the list of [[SubqueryExpression]]s in the context of current operator.
   */
  def resolveScalarSubquery(unresolvedScalarSubquery: ScalarSubquery): Expression = {
    traversals.current.parentOperator match {
      case unresolvedInlineTable: UnresolvedInlineTable =>
        throw QueryCompilationErrors.inlineTableContainsScalarSubquery(unresolvedInlineTable)
      case _ =>
    }

    val resolvedSubqueryExpressionPlan = resolveSubqueryExpressionPlan(
      unresolvedScalarSubquery.plan
    )

    val resolvedScalarSubquery = unresolvedScalarSubquery.copy(
      plan = resolvedSubqueryExpressionPlan.plan,
      outerAttrs = resolvedSubqueryExpressionPlan.outerExpressions
    )

    val coercedScalarSubquery = coerceExpressionTypes(
      expression = resolvedScalarSubquery,
      expressionTreeTraversal = traversals.current
    ).asInstanceOf[ScalarSubquery]

    expressionResolutionContextStack.peek().hasCorrelatedScalarSubqueryExpressions =
      coercedScalarSubquery.outerAttrs.nonEmpty

    operatorResolutionContextStack.current.subqueryExpressionsToValidate.add(coercedScalarSubquery)

    coercedScalarSubquery
  }

  /**
   * Resolve [[InSubquery]]:
   *  - Resolve the underlying [[ListQuery]];
   *  - Resolve the values;
   *  - Type coerce it;
   *  - Add it to the list of [[SubqueryExpression]]s in the context of current operator.
   */
  def resolveInSubquery(unresolvedInSubquery: InSubquery): Expression = {
    val resolvedQuery =
      expressionResolver.resolve(unresolvedInSubquery.query).asInstanceOf[ListQuery]

    val resolvedValues = unresolvedInSubquery.values.map { value =>
      expressionResolver.resolve(value)
    }

    val resolvedInSubquery =
      unresolvedInSubquery.copy(values = resolvedValues, query = resolvedQuery)

    val coercedInSubquery = coerceExpressionTypes(
      expression = resolvedInSubquery,
      expressionTreeTraversal = traversals.current
    ).asInstanceOf[InSubquery]

    operatorResolutionContextStack.current.subqueryExpressionsToValidate
      .add(coercedInSubquery.query)

    coercedInSubquery
  }

  /**
   * Resolve [[ListSubquery]], which is always a child of the [[InSubquery]].
   */
  def resolveListQuery(unresolvedListQuery: ListQuery): Expression = {
    val resolvedSubqueryExpressionPlan = resolveSubqueryExpressionPlan(unresolvedListQuery.plan)

    unresolvedListQuery.copy(
      plan = resolvedSubqueryExpressionPlan.plan,
      outerAttrs = resolvedSubqueryExpressionPlan.outerExpressions,
      numCols = resolvedSubqueryExpressionPlan.output.size
    )
  }

  /**
   * Resolve [[Exists]] subquery:
   *  - Resolve the subquery plan;
   *  - Get outer references;
   *  - Type coerce it;
   *  - Add it to the list of [[SubqueryExpression]]s in the context of current operator.
   */
  def resolveExists(unresolvedExists: Exists): Expression = {
    val resolvedSubqueryExpressionPlan = resolveSubqueryExpressionPlan(unresolvedExists.plan)

    val resolvedExists = unresolvedExists.copy(
      plan = resolvedSubqueryExpressionPlan.plan,
      outerAttrs = resolvedSubqueryExpressionPlan.outerExpressions
    )

    val coercedExists = coerceExpressionTypes(
      expression = resolvedExists,
      expressionTreeTraversal = traversals.current
    ).asInstanceOf[Exists]

    operatorResolutionContextStack.current.subqueryExpressionsToValidate.add(coercedExists)

    coercedExists
  }

  /**
   * Resolve [[SubqueryExpression]] plan. For the outer expressions we reuse
   * [[SubExprUtils.getOuterReferences]] for the top-down traversal to stay compatible with the
   * fixed-point Analyzer, because the order of [[SubqueryExpression.outerAttrs]] is a part of an
   * implicit alias name for that expression. Subquery expressions require:
   *  - Fresh [[NameScope]] to isolate the name resolution;
   *  - Fresh [[ExpressionIdAssigner]] mapping, because it's a separate plan branch;
   *  - Fresh [[CteScope]] with the root flag set so that [[CTERelationDefs]] are merged
   *    under the subquery root.
   *  - Fresh [[SubqueryRegistry]] scope to isolate the subquery expression plan resolution.
   *  - Fresh [[OperatorResolutionContext]] in order to preserve the context of the outer operator.
   * Collect outer expressions because they may be required for [[Window]] resolution. For
   * further context look at [[CollectsWindowSourceExpressions]].
   */
  private def resolveSubqueryExpressionPlan(
      unresolvedSubqueryPlan: LogicalPlan): ResolvedSubqueryExpressionPlan = {
    val parentSubqueryScope = subqueryRegistry.currentScope
    expressionIdAssigner.pushMapping(isSubqueryRoot = true)
    scopes.pushScope(isSubqueryRoot = true)
    cteRegistry.pushScope(isRoot = true)
    subqueryRegistry.pushScope(
      parentOperator = traversals.current.parentOperator,
      autoGeneratedAliasProvider = autoGeneratedAliasProvider,
      outerOperatorResolutionContext = operatorResolutionContextStack.current
    )
    operatorResolutionContextStack.push(unresolvedSubqueryPlan, isSubqueryRoot = true)

    val resolvedSubqueryExpressionPlan = try {
      val resolvedPlan = resolver.resolve(unresolvedSubqueryPlan)

      val outerReferences = processOuterReferences(resolvedPlan, parentSubqueryScope)

      ResolvedSubqueryExpressionPlan(
        plan = resolvedPlan,
        output = scopes.current.output,
        outerExpressions = outerReferences
      )
    } finally {
      operatorResolutionContextStack.pop()
      subqueryRegistry.popScope()
      cteRegistry.popScope()
      scopes.popScope()
      expressionIdAssigner.popMapping()
    }

    for (expression <- resolvedSubqueryExpressionPlan.outerExpressions) {
      expressionResolver.validateExpressionUnderSupportedOperator(expression)
    }

    resolvedSubqueryExpressionPlan
  }

  /**
   * Return `outerAttrs` for the resolved subquery plan.
   * If `spark.sql.optimizer.supportNestedCorrelatedSubqueries.enabled` is set to false, simply
   * collect the outer references within the subquery plan.
   *
   * Otherwise,
   *  1. Get the outer references within the resolved subquery plan and also get the outer
   *  attributes from any nested subqueries.
   *  2. For each outer reference, using the `outerReferenceScopeLevel` map from current subquery
   *  scope to check if it has only one level of correlation or beyond one level of correlation.
   *  If it is beyond one level of correlation, wrap it with [[OuterScopeReference]] and add it to
   *  the parent subquery scope's `outerReferenceScopeLevel` map with level - 1. E.g: For the
   *  following query:
   *
   * {{{
   *    SELECT *
   *    FROM t1
   *    WHERE t1.a = (
   *      SELECT MAX(t2.a)
   *      FROM t2
   *      WHERE t2.a = (
   *        SELECT MAX(t3.a)
   *        FROM t3
   *        WHERE t3.b > t2.b AND t3.c > t1.c
   *   ) AND t2.b > t1.b
   * }}}
   *
   * After resolving:
   *
   * {{{
   *   Aggregate (MAX(t3.a))
   *     Filter (t3.b > OuterReference(t2.b) AND t3.c > OuterReference(t1.c))
   *       Scan t3
   * }}}
   *
   * The innermost subquery has two outer references: `t2.b(outerScopeReferenceLevel = 1)` and
   * `t1.c(outerScopeReferenceLevel = 2)`. The innermost subquery expression is resolved as
   * `scalar-subquery(outerAttrs = [t2.b, OuterScopeReference(t1.c)])` and
   * `(t1.c.exprId -> outerScopeReferenceLevel = 1)` is added to the `outerReferenceScopeLevel`
   * map of the parent subquery scope.
   * After resolving the parent subquery plan:
   *
   * {{{
   *    Aggregate (MAX(t2.a))
   *     Filter (t2.a = scalar-subquery(outerAttrs = [t2.b, OuterScopeReference(t1.c)]) AND
   *              t2.b > OuterReference(t1.b))
   *       Scan t2
   * }}}
   *
   * The collected outer references are `t1.b` and the `OuterScopeReference(t1.c)` from the
   * innermost subquery. Both of them have `outerScopeReferenceLevel = 1`, they are all resolved
   * as one level outer references. The subquery expression is then resolved as
   * `scalar-subquery(outerAttrs = [t1.b, t1.c])`.
   */
  private def processOuterReferences(
      resolvedPlan: LogicalPlan,
      parentSubqueryScope: SubqueryScope): Seq[Expression] = {
    SubExprUtils.getOuterReferences(resolvedPlan)
  }

}
