/*
 * 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.BinaryArithmeticWithDatetimeResolver
import org.apache.spark.sql.catalyst.expressions.{
  Add,
  BinaryArithmetic,
  Cast,
  Divide,
  Expression,
  Multiply,
  Subtract
}
import org.apache.spark.sql.types._

/**
 * [[BinaryArithmeticResolver]] is invoked by [[ExpressionResolver]] in order to resolve
 * [[BinaryArithmetic]] nodes. During resolution, calling [[BinaryArithmeticWithDatetimeResolver]]
 * and applying type coercion can result in [[BinaryArithmetic]] producing some other type of node
 * or a subtree of nodes. In such cases a downwards traversal is necessary, but not going deeper
 * than the original expression's children, since all nodes below that point are guaranteed to be
 * already resolved.
 *
 * For example, given a query:
 *
 *  SELECT '4 11:11' - INTERVAL '4 22:12' DAY TO MINUTE
 *
 * [[BinaryArithmeticResolver]] is called for the following expression:
 *
 *     Subtract(
 *         Literal('4 11:11', StringType),
 *         Literal(Interval('4 22:12' DAY TO MINUTE), DayTimeIntervalType(0,2))
 *     )
 *
 * After calling [[BinaryArithmeticWithDatetimeResolver]] and applying type coercion,
 * the expression is transformed into:
 *
 *     Cast(
 *         DatetimeSub(
 *             TimestampAddInterval(
 *                 Literal('4 11:11', StringType),
 *                 UnaryMinus(
 *                     Literal(Interval('4 22:12' DAY TO MINUTE), DayTimeIntervalType(0,2))
 *                 )
 *             )
 *         )
 *     )
 *
 * A single [[Subtract]] node is replaced with a subtree of nodes. In order to resolve this subtree
 * we need to invoke [[ExpressionResolver]] recursively on the top-most node's children. The
 * top-most node itself is not resolved recursively in order to avoid recursive calls to
 * [[BinaryArithmeticResolver]] and other sub-resolvers. To prevent a case where we resolve the
 * same node twice, we need to mark nodes that will act as a limit for the downwards traversal by
 * applying a [[ResolverTag.SINGLE_PASS_SUBTREE_BOUNDARY]] tag to them. These children along with
 * all the nodes below them are guaranteed to be resolved at this point. When
 * [[ExpressionResolver]] reaches one of the tagged nodes, it returns identity rather than
 * resolving it. Finally, after resolving the subtree, we need to resolve the top-most node itself,
 * which in this case means applying a timezone, if necessary.
 */
class BinaryArithmeticResolver(expressionResolver: ExpressionResolver)
    extends TreeNodeResolver[BinaryArithmetic, Expression]
    with ProducesUnresolvedSubtree
    with CoercesExpressionTypes {

  private val traversals = expressionResolver.getExpressionTreeTraversals

  override def resolve(unresolvedBinaryArithmetic: BinaryArithmetic): Expression = {
    val binaryArithmeticWithResolvedChildren: BinaryArithmetic =
      withResolvedChildren(unresolvedBinaryArithmetic, expressionResolver.resolve _)
        .asInstanceOf[BinaryArithmetic]

    val binaryArithmeticWithResolvedSubtree: Expression =
      withResolvedSubtree(binaryArithmeticWithResolvedChildren, expressionResolver.resolve) {
        transformBinaryArithmeticNode(binaryArithmeticWithResolvedChildren)
      }

    TimezoneAwareExpressionResolver.resolveTimezone(
      binaryArithmeticWithResolvedSubtree,
      traversals.current.sessionLocalTimeZone
    )
  }

  /**
   * Transform [[BinaryArithmetic]] node by calling [[BinaryArithmeticWithDatetimeResolver]] and
   * applying type coercion. Initial node can be replaced with some other type of node or a subtree
   * of nodes.
   *
   * [[replaceDataType]] needs to be called twice, because [[BinaryArithmeticWithDatetimeResolver]]
   * algorithm is two-pass - it has separate pattern matches for the left sides, and for the right
   * sides. For example both `Subtract(_: DateType, _: StringType)` and
   * `Subtract(_: DateType, _: StringType)` are being transformed into `SubtractDates`. Same thing
   * is for the `TimeType` (`SubtractTimes`).
   */
  private def transformBinaryArithmeticNode(binaryArithmetic: BinaryArithmetic): Expression = {
    val binaryArithmeticWithNullReplaced: Expression = replaceNullType(binaryArithmetic)

    val binaryArithmeticWithDateTypeReplaced: Expression =
      replaceDateType(binaryArithmeticWithNullReplaced)

    val binaryArithmeticWithTypeCoercion: Expression =
      coerceExpressionTypes(
        expression = binaryArithmeticWithDateTypeReplaced,
        expressionTreeTraversal = traversals.current
      )

    replaceDateType(binaryArithmeticWithTypeCoercion)
  }

  /**
   * When DateType like operand is given to [[BinaryArithmetic]], apply
   * [[BinaryArithmeticWithDatetimeResolver]] in order to replace the [[BinaryArithmetic]] with
   * the appropriate equivalent for DateTime types.
   */
  private def replaceDateType(expression: Expression) = expression match {
    case arithmetic @ (_: Add | _: Subtract | _: Multiply | _: Divide) =>
      BinaryArithmeticWithDatetimeResolver.resolve(arithmetic)
    case other => other
  }

  /**
   * Replaces NullType by a compatible type in arithmetic expressions over Datetime operands.
   * This avoids recursive calls of [[BinaryArithmeticWithDatetimeResolver]] which converts
   * unacceptable nulls of `NullType` to an expected types of datetime expressions at the
   * first step, and replacing arithmetic `Add` and `Subtract` by the same datetime expressions
   * on the following steps.
   */
  private def replaceNullType(expression: Expression): Expression = expression match {
    case a @ Add(l, r, _) =>
      (l.dataType, r.dataType) match {
        case (_: DatetimeType, _: NullType) =>
          a.copy(right = Cast(a.right, DayTimeIntervalType.DEFAULT))
        case (_: NullType, _: DatetimeType) =>
          a.copy(left = Cast(a.left, DayTimeIntervalType.DEFAULT))
        case _ => a
      }
    case s @ Subtract(l, r, _) =>
      (l.dataType, r.dataType) match {
        case (_: NullType, _: DatetimeType) =>
          s.copy(left = Cast(s.left, s.right.dataType))
        case (_: DatetimeType, _: NullType) =>
          s.copy(right = Cast(s.right, s.left.dataType))
        case _ => s
      }
    case other => other
  }
}
