/*
 * 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.ignite

import org.apache.ignite.IgniteException
import org.apache.ignite.spark.impl.optimization._
import org.apache.ignite.spark.impl.optimization.accumulator._
import org.apache.ignite.spark.impl.{IgniteSQLAccumulatorRelation, IgniteSQLRelation, sqlCacheName}
import org.apache.spark.internal.Logging
import org.apache.spark.sql.catalyst.expressions._
import org.apache.spark.sql.catalyst.plans.logical._
import org.apache.spark.sql.catalyst.rules.Rule
import org.apache.spark.sql.execution.datasources.LogicalRelation

/**
  * Query plan optimization for a Ignite based queries.
  */
object IgniteOptimization extends Rule[LogicalPlan] with Logging {
    /** @inheritdoc */
    override def apply(plan: LogicalPlan): LogicalPlan = {
        logDebug("")
        logDebug("== Plan Before Ignite Operator Push Down ==")
        logDebug(plan.toString())

        val transformed = fixAmbiguousOutput(pushDownOperators(plan))

        logDebug("")
        logDebug("== Plan After Ignite Operator Push Down ==")
        logDebug(transformed.toString())

        makeIgniteAccRelation(transformed)
    }

    /**
      * Change query plan by accumulating query parts supported by Ignite into `QueryAccumulator`.
      *
      * @param plan Query plan.
      * @return Transformed plan.
      */
    private def pushDownOperators(plan: LogicalPlan): LogicalPlan = {
        val aliasIndexIterator = Stream.from(1).iterator

        //Flag to indicate that some step was skipped due to unsupported expression.
        //When it true we has to skip entire transformation of higher level Nodes.
        var stepSkipped = true

        //Applying optimization rules from bottom to up tree nodes.
        plan.transformUp {
            //We found basic node to transform.
            //We create new accumulator and going to the upper layers.
            case LogicalRelation(igniteSqlRelation: IgniteSQLRelation[_, _], output, _catalogTable, _) ⇒
                //Clear flag to optimize each statement separately
                stepSkipped = false

                val igniteQueryContext = IgniteQueryContext(
                    igniteContext = igniteSqlRelation.ic,
                    sqlContext = igniteSqlRelation.sqlContext,
                    catalogTable = _catalogTable,
                    aliasIndex = aliasIndexIterator,
                    cacheName =
                        sqlCacheName(igniteSqlRelation.ic.ignite(), igniteSqlRelation.tableName,
                            igniteSqlRelation.schemaName)
                            .getOrElse(throw new IgniteException("Unknown table")))

                //Logical Relation is bottomest TreeNode in LogicalPlan.
                //We replace it with accumulator.
                //We push all supported SQL operator into it on the higher tree levels.
                SingleTableSQLAccumulator(
                    igniteQueryContext = igniteQueryContext,
                    table = Some(igniteSqlRelation.tableName),
                    tableExpression = None,
                    outputExpressions = output.map(attr ⇒ attr.withQualifier(Seq(igniteSqlRelation.tableName))))

            case project: Project if !stepSkipped && exprsAllowed(project.projectList) ⇒
                //Project layer just changes output of current query.
                project.child match {
                    case acc: SelectAccumulator ⇒
                        acc.withOutputExpressions(
                            substituteExpressions(project.projectList, acc.outputExpressions))

                    case _ ⇒
                        throw new IgniteException("stepSkipped == true but child is not SelectAccumulator")
                }

            case sort: Sort if !stepSkipped && isSortPushDownAllowed(sort.order, sort.global) ⇒
                sort.child match {
                    case acc: QueryAccumulator ⇒
                        acc.withOrderBy(sort.order)

                    case _ ⇒
                        throw new IgniteException("stepSkipped == true but child is not SelectAccumulator")
                }

            case filter: Filter if !stepSkipped && exprsAllowed(filter.condition) ⇒

                filter.child match {
                    case acc: SelectAccumulator ⇒
                        if (hasAggregateInside(filter.condition) || acc.groupBy.isDefined)
                            acc.withHaving(acc.having.getOrElse(Nil) :+ filter.condition)
                        else
                            acc.withWhere(acc.where.getOrElse(Nil) :+ filter.condition)

                    case _ ⇒
                        throw new IgniteException("stepSkipped == true but child is not SelectAccumulator")
                }

            case agg: Aggregate
                if !stepSkipped && exprsAllowed(agg.groupingExpressions) && exprsAllowed(agg.aggregateExpressions) ⇒

                agg.child match {
                    case acc: SelectAccumulator ⇒
                        if (acc.groupBy.isDefined) {
                            val tableAlias = acc.igniteQueryContext.uniqueTableAlias

                            SingleTableSQLAccumulator(
                                igniteQueryContext = acc.igniteQueryContext,
                                table = None,
                                tableExpression = Some((acc, tableAlias)),
                                outputExpressions = agg.aggregateExpressions)
                        }
                        else
                            acc
                                .withGroupBy(agg.groupingExpressions)
                                .withOutputExpressions(
                                    substituteExpressions(agg.aggregateExpressions, acc.outputExpressions))

                    case acc: QueryAccumulator ⇒
                        val tableAlias = acc.igniteQueryContext.uniqueTableAlias

                        SingleTableSQLAccumulator(
                            igniteQueryContext = acc.igniteQueryContext,
                            table = None,
                            tableExpression = Some((acc, tableAlias)),
                            outputExpressions = agg.aggregateExpressions)

                    case _ ⇒
                        throw new IgniteException("stepSkipped == true but child is not SelectAccumulator")
                }

            case limit: LocalLimit if !stepSkipped && exprsAllowed(limit.limitExpr) ⇒
                limit.child match {
                    case acc: SelectAccumulator ⇒
                        acc.withLocalLimit(limit.limitExpr)

                    case acc: QueryAccumulator ⇒
                        acc.withLocalLimit(limit.limitExpr)

                    case _ ⇒
                        throw new IgniteException("stepSkipped == true but child is not SelectAccumulator")
                }

            case limit: GlobalLimit if !stepSkipped && exprsAllowed(limit.limitExpr) ⇒
                limit.child.transformUp {
                    case acc: SelectAccumulator ⇒
                        acc.withLimit(limit.limitExpr)

                    case acc: QueryAccumulator ⇒
                        acc.withLimit(limit.limitExpr)

                    case _ ⇒
                        throw new IgniteException("stepSkipped == true but child is not SelectAccumulator")
                }

            case union: Union if !stepSkipped && isAllChildrenOptimized(union.children) ⇒
                val first = union.children.head.asInstanceOf[QueryAccumulator]

                val subQueries = union.children.map(_.asInstanceOf[QueryAccumulator])

                UnionSQLAccumulator(
                    first.igniteQueryContext,
                    subQueries,
                    subQueries.head.output)

            case join: Join
                if !stepSkipped && isAllChildrenOptimized(Seq(join.left, join.right)) &&
                    join.condition.forall(exprsAllowed) ⇒

                val left = join.left.asInstanceOf[QueryAccumulator]

                val (leftOutput, leftAlias) =
                    if (!isSimpleTableAcc(left)) {
                        val tableAlias = left.igniteQueryContext.uniqueTableAlias

                        (left.output, Some(tableAlias))
                    }
                    else
                        (left.output, None)

                val right = join.right.asInstanceOf[QueryAccumulator]

                val (rightOutput, rightAlias) =
                    if (!isSimpleTableAcc(right) ||
                        leftAlias.getOrElse(left.qualifier) == right.qualifier) {
                        val tableAlias = right.igniteQueryContext.uniqueTableAlias

                        (right.output, Some(tableAlias))
                    }
                    else
                        (right.output, None)

                JoinSQLAccumulator(
                    left.igniteQueryContext,
                    left,
                    right,
                    join.joinType,
                    leftOutput ++ rightOutput,
                    join.condition,
                    leftAlias,
                    rightAlias)

            case unknown ⇒
                stepSkipped = true

                unknown
        }
    }

    /**
      * Changes qualifiers for an ambiguous columns names.
      *
      * @param plan Query plan.
      * @return Transformed plan.
      */
    private def fixAmbiguousOutput(plan: LogicalPlan): LogicalPlan = plan.transformDown {
        case acc: SingleTableSQLAccumulator if acc.children.exists(_.isInstanceOf[JoinSQLAccumulator]) ⇒
            val fixedChildOutput =
                fixAmbiguousOutput(acc.tableExpression.get._1.outputExpressions, acc.igniteQueryContext)

            val newOutput = substituteExpressions(acc.outputExpressions, fixedChildOutput, changeOnlyName = true)

            acc.copy(
                outputExpressions = newOutput,
                where = acc.where.map(
                    substituteExpressions(_, fixedChildOutput, changeOnlyName = true)),
                groupBy = acc.groupBy.map(
                    substituteExpressions(_, fixedChildOutput, changeOnlyName = true)),
                having = acc.having.map(
                    substituteExpressions(_, fixedChildOutput, changeOnlyName = true)),
                limit = acc.limit.map(
                    substituteExpression(_, fixedChildOutput, changeOnlyName = true)),
                localLimit = acc.localLimit.map(
                    substituteExpression(_, fixedChildOutput, changeOnlyName = true)),
                orderBy = acc.orderBy.map(
                    substituteExpressions(_, fixedChildOutput, changeOnlyName = true)))

            acc

        case acc: JoinSQLAccumulator
            if acc.left.isInstanceOf[JoinSQLAccumulator] || acc.right.isInstanceOf[JoinSQLAccumulator] ⇒
            val leftFixed = acc.left match {
                case leftJoin: JoinSQLAccumulator ⇒
                    val fixedChildOutput = fixAmbiguousOutput(acc.left.outputExpressions, acc.igniteQueryContext)

                    val newOutput =
                        substituteExpressions(acc.outputExpressions, fixedChildOutput, changeOnlyName = true)

                    acc.copy(
                        outputExpressions = newOutput,
                        left = leftJoin.copy(outputExpressions = fixedChildOutput),
                        condition = acc.condition.map(
                            substituteExpression(_, fixedChildOutput, changeOnlyName = true)),
                        where = acc.where.map(
                            substituteExpressions(_, fixedChildOutput, changeOnlyName = true)),
                        groupBy = acc.groupBy.map(
                            substituteExpressions(_, fixedChildOutput, changeOnlyName = true)),
                        having = acc.having.map(
                            substituteExpressions(_, fixedChildOutput, changeOnlyName = true)),
                        limit = acc.limit.map(
                            substituteExpression(_, fixedChildOutput, changeOnlyName = true)),
                        localLimit = acc.localLimit.map(
                            substituteExpression(_, fixedChildOutput, changeOnlyName = true)),
                        orderBy = acc.orderBy.map(
                            substituteExpressions(_, fixedChildOutput, changeOnlyName = true)))

                case _ ⇒ acc
            }

            val fixed = leftFixed.right match {
                case rightJoin: JoinSQLAccumulator ⇒
                    val fixedChildOutput =
                        fixAmbiguousOutput(leftFixed.outputExpressions, leftFixed.igniteQueryContext)

                    val newOutput = substituteExpressions(leftFixed.outputExpressions, fixedChildOutput)

                    leftFixed.copy(
                        outputExpressions = newOutput,
                        right = rightJoin.copy(outputExpressions = fixedChildOutput),
                        condition = acc.condition.map(
                            substituteExpression(_, fixedChildOutput, changeOnlyName = true)),
                        where = acc.where.map(
                            substituteExpressions(_, fixedChildOutput, changeOnlyName = true)),
                        groupBy = acc.groupBy.map(
                            substituteExpressions(_, fixedChildOutput, changeOnlyName = true)),
                        having = acc.having.map(
                            substituteExpressions(_, fixedChildOutput, changeOnlyName = true)),
                        limit = acc.limit.map(
                            substituteExpression(_, fixedChildOutput, changeOnlyName = true)),
                        localLimit = acc.localLimit.map(
                            substituteExpression(_, fixedChildOutput, changeOnlyName = true)),
                        orderBy = acc.orderBy.map(
                            substituteExpressions(_, fixedChildOutput, changeOnlyName = true)))

                case _ ⇒ leftFixed
            }

            fixed.copy(
                condition = acc.condition.map(
                    substituteExpression(_, acc.outputExpressions, changeOnlyName = true)),
                where = acc.where.map(
                    substituteExpressions(_, acc.outputExpressions, changeOnlyName = true)),
                groupBy = acc.groupBy.map(
                    substituteExpressions(_, acc.outputExpressions, changeOnlyName = true)),
                having = acc.having.map(
                    substituteExpressions(_, acc.outputExpressions, changeOnlyName = true)),
                limit = acc.limit.map(
                    substituteExpression(_, acc.outputExpressions, changeOnlyName = true)),
                localLimit = acc.localLimit.map(
                    substituteExpression(_, acc.outputExpressions, changeOnlyName = true)),
                orderBy = acc.orderBy.map(
                    substituteExpressions(_, acc.outputExpressions, changeOnlyName = true)))

        case unknown ⇒
            unknown
    }

    private def fixAmbiguousOutput(exprs: Seq[NamedExpression], ctx: IgniteQueryContext): Seq[NamedExpression] =
        exprs.foldLeft((Set[String](), Set[NamedExpression]())) {
            case ((uniqueNames, fixed), next) ⇒
                if (uniqueNames(next.name))
                    (uniqueNames, fixed + Alias(next, ctx.uniqueColumnAlias(next))(exprId = next.exprId))
                else
                    (uniqueNames + next.name, fixed + next)
        }._2.toSeq

    /**
      * Substitutes each `QueryAccumulator` with a `LogicalRelation` contains `IgniteSQLAccumulatorRelation`.
      *
      * @param plan Query plan.
      * @return Transformed plan.
      */
    private def makeIgniteAccRelation(plan: LogicalPlan): LogicalPlan =
        plan.transformDown {
            case acc: QueryAccumulator ⇒
                new LogicalRelation (
                    relation = IgniteSQLAccumulatorRelation(acc),
                    output = acc.outputExpressions.map(toAttributeReference(_, Seq.empty)),
                    catalogTable = acc.igniteQueryContext.catalogTable,
                    false)
        }

    /**
      * @param order Order.
      * @param global True if order applied to entire result set false if ordering per-partition.
      * @return True if sort can be pushed down to Ignite, false otherwise.
      */
    private def isSortPushDownAllowed(order: Seq[SortOrder], global: Boolean): Boolean =
        global && order.map(_.child).forall(exprsAllowed)

    /**
      * @param children Plans to check.
      * @return True is all plan are `QueryAccumulator`, false otherwise.
      */
    private def isAllChildrenOptimized(children: Seq[LogicalPlan]): Boolean =
        children.forall {
            case _: QueryAccumulator ⇒
                true

            case _ ⇒
                false
        }

    /**
      * Changes expression from `exprs` collection to expression with same `exprId` from `substitution`.
      *
      * @param exprs Expressions to substitute.
      * @param substitution Expressions for substitution
      * @param changeOnlyName If true substitute only expression name.
      * @tparam T Concrete expression type.
      * @return Substituted expressions.
      */
    private def substituteExpressions[T <: Expression](exprs: Seq[T], substitution: Seq[NamedExpression],
        changeOnlyName: Boolean = false): Seq[T] = {

        exprs.map(substituteExpression(_, substitution, changeOnlyName))
    }

    private def substituteExpression[T <: Expression](expr: T, substitution: Seq[NamedExpression],
        changeOnlyName: Boolean): T = expr match {
        case ne: NamedExpression ⇒
            substitution.find(_.exprId == ne.exprId) match {
                case Some(found) ⇒
                    if (!changeOnlyName)
                        found.asInstanceOf[T]
                    else ne match {
                        case alias: Alias ⇒
                            Alias(
                                AttributeReference(
                                    found.name,
                                    found.dataType,
                                    nullable = found.nullable,
                                    metadata = found.metadata)(
                                    exprId = found.exprId,
                                    qualifier = found.qualifier),
                                alias.name) (
                                exprId = alias.exprId,
                                qualifier = alias.qualifier,
                                explicitMetadata = alias.explicitMetadata).asInstanceOf[T]

                        case attr: AttributeReference ⇒
                            attr.copy(name = found.name)(
                                exprId = found.exprId,
                                qualifier = found.qualifier).asInstanceOf[T]

                        case _ ⇒ ne.asInstanceOf[T]
                    }

                case None ⇒
                    expr.withNewChildren(
                        substituteExpressions(expr.children, substitution, changeOnlyName)).asInstanceOf[T]
            }

        case _ ⇒
            expr.withNewChildren(
                substituteExpressions(expr.children, substitution, changeOnlyName)).asInstanceOf[T]
    }
}
