package com.foo

import java.util

import com.google.common.collect.ImmutableList
import org.apache.calcite.plan.hep.HepRelVertex
import org.apache.calcite.plan.{RelOptRule, RelOptRuleCall, RelOptRuleOperand}
import org.apache.calcite.rel.core.AggregateCall
import org.apache.calcite.rel.logical.{LogicalAggregate, LogicalProject}
import org.apache.calcite.rex.RexNode
import org.apache.calcite.util.ImmutableBitSet

import scala.collection.JavaConverters._

/*
== BEFORE
LogicalAggregate(group=[{0}], EXPR$1=[SUM($1)])
  LogicalProject(id=[$0], $f1=[CASE(=($1, 'foo'), $2, 0)])
    LogicalAggregate(group=[{0, 1}], sum_units=[SUM($2)])
      LogicalTableScan(table=[[SALES, Orders]])
*/

/*
== AFTER
LogicalAggregate(group=[{0}], EXPR$1=[SUM($1)])
  LogicalProject(id=[$0], $f1=[$1])
    LogicalAggregate(group=[{0}], sum_units=[SUM($1)])
      LogicalProject(id=[$0], sum_units=[CASE(=($1, 'foo'), $2, 0)])
        LogicalTableScan(table=[[SALES, Orders]])

 */

class PushDownRule(operand: RelOptRuleOperand) extends RelOptRule(operand) {

  override def onMatch(call: RelOptRuleCall): Unit = {

    val rexBuilder = call.builder().getRexBuilder
    val rootLogProj = call.rels(0).asInstanceOf[LogicalProject]

    if (call.id == 0) {

      val logAgg = rootLogProj.getInput.asInstanceOf[HepRelVertex].getCurrentRel.asInstanceOf[LogicalAggregate]
      val tblScan = logAgg.getInput

      // create new logical project to be inserted before table scan
      val rexNodes = new java.util.ArrayList[RexNode]()

      val idRef = rexBuilder.makeInputRef(tblScan, 0)
      // val productRef = rexBuilder.makeInputRef(tblScan, 1)
      val exprRef = rootLogProj.getChildExps.get(1)

      rexNodes.add(idRef)
      // rexNodes.add(productRef)
      rexNodes.add(exprRef)

      val rexNames = new java.util.ArrayList[String]()
      rexNames.add("id")
      // rexNames.add("product")
      rexNames.add("sum_units")

      val newLogProj = LogicalProject.create(logAgg.getInput, rexNodes, rexNames)

      // recreate L2 aggregate with new logical project
      val groups = new util.ArrayList[Integer]
      groups.add(0)
      val newLogAggGroupSet = ImmutableBitSet.of(groups)
      val newLogAggGroupSets = ImmutableList.of(newLogAggGroupSet)

      val newAggCallList = new util.ArrayList[AggregateCall]()
      val newIndexes = new util.ArrayList[Integer]()
      newIndexes.add(1)
      val newAggCall = logAgg.getAggCallList.get(0).copy(newIndexes, -1)
      newAggCallList.add(newAggCall)

      val newLogAgg = logAgg.copy(logAgg.getTraitSet, newLogProj, logAgg.indicator, newLogAggGroupSet, newLogAggGroupSets, newAggCallList)

      // recreate L1 logical project with the new expressions and parent aggregate
      val newRootLogProjRexNodes = new java.util.ArrayList[RexNode]()

      val rootLogProjIdRef = rexBuilder.makeInputRef(newLogAgg, 0)
      val rootLogProjUnitsRef = rexBuilder.makeInputRef(newLogAgg, 1)

      newRootLogProjRexNodes.add(rootLogProjIdRef)
      newRootLogProjRexNodes.add(rootLogProjUnitsRef)

      val newRootLogProj = rootLogProj.copy(rootLogProj.getTraitSet, newLogAgg, newRootLogProjRexNodes, rootLogProj.getRowType)

      call.transformTo(newRootLogProj)
    }
  }
}

object PushDownRule {
  def apply: PushDownRule = {
    import RelOptRule._

    new PushDownRule(
      operand(classOf[LogicalProject], any()))
  }
}