/*
 * 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.PivotTransformer
import org.apache.spark.sql.catalyst.expressions.{
  AliasHelper,
  Attribute,
  Expression,
  NamedExpression,
  RowOrdering
}
import org.apache.spark.sql.catalyst.plans.logical.{Aggregate, LogicalPlan, Pivot, Project}
import org.apache.spark.sql.errors.QueryCompilationErrors

/**
 * Resolver class that resolves [[Pivot]] operators.
 */
class PivotResolver(operatorResolver: Resolver, expressionResolver: ExpressionResolver)
    extends TreeNodeResolver[Pivot, LogicalPlan]
    with AliasHelper
    with CoercesExpressionTypes {
  private val scopes: NameScopeStack = operatorResolver.getNameScopes
  private val autoGeneratedAliasProvider = expressionResolver.getAutoGeneratedAliasProvider

  /**
   * Resolves a [[Pivot]] operator by:
   *  1. Resolving its child operator.
   *  2. Resolving its group by expressions (if any).
   *  3. Resolving its pivot column.
   *  4. Resolving its pivot values.
   *  5. Resolving its aggregate expressions.
   *  6. Validating that the pivot column is orderable.
   *  7. Constructing the resolved operator using [[PivotTransformer]] (see its doc for details).
   *  8. Handling the resolved operator to resolve any remaining expressions.
   *  9. Updating the current scope with the new output attributes.
   *
   * For example consider the following query:
   *
   * {{{
   *   SELECT * FROM VALUES (2024, 1) table(c1, c2) PIVOT (sum(c1) FOR c2 IN (1 AS q1));
   * }}}
   *
   * The parsed plan would be:
   *
   * {{{
   *   'Project [*]
   *   +- 'Pivot 'c2, [1 AS q1#2], ['sum('c1)]
   *      +- SubqueryAlias table
   *         +- LocalRelation [c1#0, c2#1]
   * }}}
   *
   * Here the `pivotColumn` is `c2`, the `pivotValues` is `[1 AS q1]`, and the `aggregates` is
   * `sum(c1)`.
   * After the resolution, the resolved plan would be:
   *
   * {{{
   *   Project [q1#4]
   *   +- Project [__pivot_sum(c1) AS `sum(c1)`#3[0] AS q1#4]
   *      +- Aggregate [pivotfirst(c2#1, sum(c1)#2, 1, 0, 0) AS __pivot_sum(c1) AS `sum(c1)`#3]
   *         +- Aggregate [c2#1], [c2#1, sum(c1#0) AS sum(c1)#2]
   *            +- LocalRelation [c1#0, c2#1]
   * }}}
   */
  override def resolve(pivot: Pivot): LogicalPlan = {
    val resolvedPivot = {
      val resolvedChild = operatorResolver.resolve(pivot.child)
      val resolvedGroupByExprsOpt = pivot.groupByExprsOpt.map { expressions =>
        expressions.map { expression =>
          expressionResolver
            .resolveExpressionTreeInOperator(
              parentOperator = pivot,
              unresolvedExpression = expression
            )
            .asInstanceOf[NamedExpression]
        }
      }
      val resolvedPivotColumn = expressionResolver.resolveExpressionTreeInOperator(
        parentOperator = pivot,
        unresolvedExpression = pivot.pivotColumn
      )
      val resolvedPivotValues = pivot.pivotValues.map { expression =>
        expressionResolver.resolveExpressionTreeInOperator(
          parentOperator = pivot,
          unresolvedExpression = expression
        )
      }
      val resolvedAggregates = expressionResolver.resolvePivotAggregates(pivot)

      checkUnorderablePivotColError(resolvedPivotColumn)

      val transformedPivot = PivotTransformer(
        child = resolvedChild,
        pivotValues = resolvedPivotValues,
        pivotColumn = resolvedPivotColumn,
        groupByExpressionsOpt = resolvedGroupByExprsOpt,
        aggregates = resolvedAggregates,
        childOutput = scopes.current.output,
        newAlias = (child, name) => autoGeneratedAliasProvider.newAlias(child, name)
      )

      handleResolvedOperator(transformedPivot)
    }

    val newOutput = computeOutput(resolvedPivot)

    scopes.overwriteOutputAndExtendHiddenOutput(output = newOutput)

    resolvedPivot
  }

  /**
   * [[PivotTransformer]] may introduce new expressions (such as [[Cast]], [[If]] etc) that need
   * to be resolved again. This method resolves those expressions in the given operator.
   */
  private def handleResolvedOperator(operator: LogicalPlan): LogicalPlan = {
    operator match {
      case project: Project =>
        val resolvedProjectList = project.projectList.map { expression =>
          expressionResolver
            .resolveExpressionTreeInOperator(
              parentOperator = project,
              unresolvedExpression = expression
            )
            .asInstanceOf[NamedExpression]
        }
        project.copy(projectList = resolvedProjectList)
      case aggregate: Aggregate =>
        val resolvedAggregateExpressions = aggregate.aggregateExpressions.map { expression =>
          expressionResolver
            .resolveExpressionTreeInOperator(
              parentOperator = aggregate,
              unresolvedExpression = expression
            )
            .asInstanceOf[NamedExpression]
        }
        aggregate.copy(aggregateExpressions = resolvedAggregateExpressions)
    }
  }

  private def computeOutput(operator: LogicalPlan): Seq[Attribute] = {
    operator match {
      case project: Project =>
        project.projectList.map(namedExpression => namedExpression.toAttribute)
      case aggregate: Aggregate =>
        aggregate.aggregateExpressions.map(namedExpression => namedExpression.toAttribute)
    }
  }

  private def checkUnorderablePivotColError(pivotColumn: Expression): Unit = {
    if (!RowOrdering.isOrderable(pivotColumn.dataType)) {
      throw QueryCompilationErrors.unorderablePivotColError(pivotColumn)
    }
  }
}
