/*
 * 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.flink.table.plan.metadata

import org.apache.flink.table.calcite.FlinkRelBuilder.NamedWindowProperty
import org.apache.flink.table.plan.metadata.FlinkMetadata.UniqueGroups
import org.apache.flink.table.plan.nodes.calcite.{Expand, LogicalWindowAggregate, Rank}
import org.apache.flink.table.plan.nodes.logical.FlinkLogicalWindowAggregate
import org.apache.flink.table.plan.nodes.physical.batch._
import org.apache.flink.table.plan.util.FlinkRelMdUtil
import org.apache.flink.table.plan.util.FlinkRelMdUtil.splitColumnsIntoLeftAndRight
import org.apache.flink.table.plan.util.FlinkRelOptUtil.checkAndSplitAggCalls

import org.apache.calcite.plan.volcano.RelSubset
import org.apache.calcite.rel.core._
import org.apache.calcite.rel.metadata._
import org.apache.calcite.rel.{RelNode, SingleRel}
import org.apache.calcite.rex._
import org.apache.calcite.sql.SqlKind
import org.apache.calcite.util.{Bug, ImmutableBitSet, Util}

import java.util

import scala.collection.JavaConversions._
import scala.collection.mutable

class FlinkRelMdUniqueGroups private extends MetadataHandler[UniqueGroups] {

  override def getDef: MetadataDef[UniqueGroups] = FlinkMetadata.UniqueGroups.DEF

  def getUniqueGroups(
      ts: TableScan,
      mq: RelMetadataQuery,
      columns: ImmutableBitSet): ImmutableBitSet = {
    val uniqueKeys = mq.getUniqueKeys(ts)
    if (uniqueKeys == null || uniqueKeys.isEmpty) {
      return columns
    }
    require(columns.forall(_ < ts.getRowType.getFieldCount))
    val none = Option.empty[ImmutableBitSet]
    // find the minimum uniqueKey
    val uniqueGroups = uniqueKeys.foldLeft(none) {
      (groups, uniqueKey) =>
        val containUniqueKey = columns.contains(uniqueKey)
        groups match {
          case Some(g) =>
            if (containUniqueKey && g.cardinality() > uniqueKey.cardinality()) {
              Some(uniqueKey)
            } else {
              groups
            }
          case _ => if (containUniqueKey) Some(uniqueKey) else none
        }
    }
    uniqueGroups.getOrElse(columns)
  }

  def getUniqueGroups(
      expand: Expand,
      mq: RelMetadataQuery,
      columns: ImmutableBitSet): ImmutableBitSet = {
    val columnList = columns.toList
    val fmq = FlinkRelMetadataQuery.reuseOrCreate(mq)
    val columnsSkipExpandId = columnList.filter(_ != expand.expandIdIndex)
    if (columnsSkipExpandId.isEmpty) {
      return columns
    }

    // mapping input column index to output index for non-null value columns
    val mapInputToOutput = new util.HashMap[Int, Int]()
    columnsSkipExpandId.foreach { column =>
      val inputRefs = FlinkRelMdUtil.getInputRefIndices(column, expand)
      if (inputRefs.size() == 1 && inputRefs.head >= 0) {
        mapInputToOutput.put(inputRefs.head, column)
      }
    }
    if (mapInputToOutput.isEmpty) {
      return columns
    }

    val leftColumns = columnList.filterNot(mapInputToOutput.values().contains)
    val inputUniqueGroups = fmq.getUniqueGroups(
      expand.getInput, ImmutableBitSet.of(mapInputToOutput.keys.toSeq: _*))
    val outputUniqueGroups = inputUniqueGroups.map(mapInputToOutput.get)
    ImmutableBitSet.of(outputUniqueGroups.toSeq: _*).union(ImmutableBitSet.of(leftColumns))
  }

  def getUniqueGroups(
      rank: Rank,
      mq: RelMetadataQuery,
      columns: ImmutableBitSet): ImmutableBitSet = {
    val columnList = columns.toList
    val fmq = FlinkRelMetadataQuery.reuseOrCreate(mq)
    val rankFunColumnIndex = FlinkRelMdUtil.getRankFunColumnIndex(rank)
    val columnSkipRankCol = columnList.filter(_ != rankFunColumnIndex)
    if (columnSkipRankCol.isEmpty) {
      return columns
    }
    val inputUniqueGroups = fmq.getUniqueGroups(
      rank.getInput, ImmutableBitSet.of(columnSkipRankCol))
    if (columnList.contains(rankFunColumnIndex)) {
      inputUniqueGroups.union(ImmutableBitSet.of(rankFunColumnIndex))
    } else {
      inputUniqueGroups
    }
  }

  def getUniqueGroups(
      filter: Filter,
      mq: RelMetadataQuery,
      columns: ImmutableBitSet): ImmutableBitSet = {
    val fmq = FlinkRelMetadataQuery.reuseOrCreate(mq)
    fmq.getUniqueGroups(filter.getInput, columns)
  }

  def getUniqueGroups(
      project: Project,
      mq: RelMetadataQuery,
      columns: ImmutableBitSet): ImmutableBitSet = {
    val projects = project.getProjects
    getUniqueGroupsOfProject(projects, project.getInput, mq, columns)
  }

  def getUniqueGroups(
      calc: Calc,
      mq: RelMetadataQuery,
      columns: ImmutableBitSet): ImmutableBitSet = {
    val projects = calc.getProgram.getProjectList.map(calc.getProgram.expandLocalRef)
    getUniqueGroupsOfProject(projects, calc.getInput, mq, columns)
  }

  private def getUniqueGroupsOfProject(
      projects: util.List[RexNode],
      input: RelNode,
      mq: RelMetadataQuery,
      columns: ImmutableBitSet): ImmutableBitSet = {
    val columnList = columns.toList
    val mapInToOutRefPos = new mutable.HashMap[Integer, Integer]()
    // find non-RexInputRef and non-Constant columns
    val outNonRefOrConstantCols = new mutable.ArrayBuffer[Integer]()
    columns.foreach { column =>
      require(column < projects.size)
      projects.get(column) match {
        case ref: RexInputRef => mapInToOutRefPos.putIfAbsent(ref.getIndex, column)
        case call: RexCall if call.getKind.equals(SqlKind.AS) &&
          call.getOperands.head.isInstanceOf[RexInputRef] =>
          val index = call.getOperands.head.asInstanceOf[RexInputRef].getIndex
          mapInToOutRefPos.putIfAbsent(index, column)
        case _: RexLiteral => // do nothing
        case _ => outNonRefOrConstantCols += column
      }
    }

    if (mapInToOutRefPos.isEmpty) {
      val nonConstantCols = columnList.filterNot { column =>
        projects.get(column).isInstanceOf[RexLiteral]
      }
      if (nonConstantCols.isEmpty) {
        // all columns are constant, return first column
        ImmutableBitSet.of(columnList.head)
      } else {
        // return non-constant columns
        ImmutableBitSet.of(nonConstantCols)
      }
    } else {
      val fmq = FlinkRelMetadataQuery.reuseOrCreate(mq)
      val inputColumns = ImmutableBitSet.of(mapInToOutRefPos.keys.toList)
      val inputUniqueGroups = fmq.getUniqueGroups(input, inputColumns)
      val outputUniqueGroups = inputUniqueGroups.asList.map {
        k => mapInToOutRefPos.getOrElse(k, throw new IllegalArgumentException(s"Illegal index: $k"))
      }
      ImmutableBitSet.of(outputUniqueGroups).union(ImmutableBitSet.of(outNonRefOrConstantCols))
    }
  }

  def getUniqueGroups(
      exchange: Exchange,
      mq: RelMetadataQuery,
      columns: ImmutableBitSet): ImmutableBitSet = {
    val fmq = FlinkRelMetadataQuery.reuseOrCreate(mq)
    fmq.getUniqueGroups(exchange.getInput, columns)
  }

  def getUniqueGroups(
      rel: SetOp,
      mq: RelMetadataQuery,
      columns: ImmutableBitSet): ImmutableBitSet = columns

  def getUniqueGroups(
      sort: Sort,
      mq: RelMetadataQuery,
      columns: ImmutableBitSet): ImmutableBitSet = {
    val fmq = FlinkRelMetadataQuery.reuseOrCreate(mq)
    fmq.getUniqueGroups(sort.getInput, columns)
  }

  def getUniqueGroups(
      rel: Correlate,
      mq: RelMetadataQuery,
      columns: ImmutableBitSet): ImmutableBitSet = columns

  def getUniqueGroups(
      rel: BatchExecCorrelate,
      mq: RelMetadataQuery,
      columns: ImmutableBitSet): ImmutableBitSet = columns

  def getUniqueGroups(
      join: Join,
      mq: RelMetadataQuery,
      columns: ImmutableBitSet): ImmutableBitSet = {
    require(join.getSystemFieldList.isEmpty)
    val leftFieldCount = join.getLeft.getRowType.getFieldCount
    val (leftColumns, rightColumns) = splitColumnsIntoLeftAndRight(leftFieldCount, columns)

    val fmq = FlinkRelMetadataQuery.reuseOrCreate(mq)
    val leftUniqueGroups = fmq.getUniqueGroups(join.getLeft, leftColumns)
    val rightUniqueGroups = fmq.getUniqueGroups(join.getRight, rightColumns)

    val joinType = join.getJoinType
    val joinInfo = join.analyzeCondition()
    val leftJoinKeys = ImmutableBitSet.of(joinInfo.leftKeys)
    val rightJoinKeys = ImmutableBitSet.of(joinInfo.rightKeys)
    // for INNER and LEFT join, returns leftUniqueGroups if the join keys of RHS are unique
    if (leftJoinKeys.nonEmpty
      && leftUniqueGroups.contains(leftJoinKeys)
      && !joinType.generatesNullsOnLeft()) {
      val isRightJoinKeysUnique = fmq.areColumnsUnique(join.getRight, rightJoinKeys)
      if (isRightJoinKeysUnique != null && isRightJoinKeysUnique) {
        return leftUniqueGroups
      }
    }

    val outputRightUniqueGroups =
      rightUniqueGroups.asList.map(c => Integer.valueOf(c + leftFieldCount))
    // for INNER and RIGHT join, returns rightUniqueGroups if the join keys of LHS are unique
    if (rightJoinKeys.nonEmpty
      && rightUniqueGroups.contains(rightJoinKeys)
      && !joinType.generatesNullsOnRight()) {
      val isLeftJoinKeysUnique = fmq.areColumnsUnique(join.getLeft, leftJoinKeys)
      if (isLeftJoinKeysUnique != null && isLeftJoinKeysUnique) {
        return ImmutableBitSet.of(outputRightUniqueGroups)
      }
    }

    leftUniqueGroups.union(ImmutableBitSet.of(outputRightUniqueGroups))
  }

  def getUniqueGroups(
      semiJoin: SemiJoin,
      mq: RelMetadataQuery,
      columns: ImmutableBitSet): ImmutableBitSet = {
    require(semiJoin.getSystemFieldList.isEmpty)
    val fmq = FlinkRelMetadataQuery.reuseOrCreate(mq)
    fmq.getUniqueGroups(semiJoin.getLeft, columns)
  }

  def getUniqueGroups(
      agg: Aggregate,
      mq: RelMetadataQuery,
      columns: ImmutableBitSet): ImmutableBitSet = {
    val grouping = agg.getGroupSet.map(_.toInt).toArray
    getUniqueGroupsOfAggregate(agg.getRowType.getFieldCount, grouping, agg.getInput, mq, columns)
  }

  def getUniqueGroups(
      agg: BatchExecGroupAggregateBase,
      mq: RelMetadataQuery,
      columns: ImmutableBitSet): ImmutableBitSet = {
    val grouping = agg.getGrouping
    getUniqueGroupsOfAggregate(agg.getRowType.getFieldCount, grouping, agg.getInput, mq, columns)
  }

  private def getUniqueGroupsOfAggregate(
      outputFiledCount: Int,
      grouping: Array[Int],
      input: RelNode,
      mq: RelMetadataQuery,
      columns: ImmutableBitSet): ImmutableBitSet = {
    val columnList = columns.toList
    val groupingInToOutMap = new mutable.HashMap[Integer, Integer]()
    columnList.foreach { column =>
      require(column < outputFiledCount)
      if (column < grouping.length) {
        groupingInToOutMap.put(grouping(column), column)
      }
    }
    if (groupingInToOutMap.isEmpty) {
      columns
    } else {
      val fmq = FlinkRelMetadataQuery.reuseOrCreate(mq)
      val inputColumns = ImmutableBitSet.of(groupingInToOutMap.keys.toList)
      val inputUniqueGroups = fmq.getUniqueGroups(input, inputColumns)
      val uniqueGroupsFromGrouping = inputUniqueGroups.asList.map { k =>
        groupingInToOutMap.getOrElse(k, throw new IllegalArgumentException(s"Illegal index: $k"))
      }
      val nonGroupingCols = if (inputColumns.toArray.sorted.sameElements(grouping.sorted)) {
        // if values of inputColumns are grouping columns, nonGroupingCols can be dropped.
        // (because grouping columns are unique.)
        Seq.empty[Integer]
      } else {
        val groupingOutColumns = groupingInToOutMap.values
        columnList.filterNot(groupingOutColumns.contains(_))
      }
      ImmutableBitSet.of(uniqueGroupsFromGrouping).union(ImmutableBitSet.of(nonGroupingCols))
    }
  }

  def getUniqueGroups(
      window: LogicalWindowAggregate,
      mq: RelMetadataQuery,
      columns: ImmutableBitSet): ImmutableBitSet = {
    val grouping = window.getGroupSet.map(_.toInt).toArray
    val namedProperties = window.getNamedProperties
    val (auxGroupSet, _) = checkAndSplitAggCalls(window)
    if (window.indicator) {
      require(auxGroupSet.isEmpty)
    }
    getUniqueGroupsOfWindow(window, grouping, auxGroupSet, namedProperties, mq, columns)
  }

  def getUniqueGroups(
      window: FlinkLogicalWindowAggregate,
      mq: RelMetadataQuery,
      columns: ImmutableBitSet): ImmutableBitSet = {
    val grouping = window.getGroupSet.map(_.toInt).toArray
    val namedProperties = window.getNamedProperties
    val (auxGroupSet, _) = checkAndSplitAggCalls(window)
    getUniqueGroupsOfWindow(window, grouping, auxGroupSet, namedProperties, mq, columns)
  }

  def getUniqueGroups(
      window: BatchExecWindowAggregateBase,
      mq: RelMetadataQuery,
      columns: ImmutableBitSet): ImmutableBitSet = {
    val grouping = window.getGrouping
    val namedProperties = window.getNamedProperties
    getUniqueGroupsOfWindow(window, grouping, window.getAuxGrouping, namedProperties, mq, columns)
  }

  private def getUniqueGroupsOfWindow(
      window: SingleRel,
      grouping: Array[Int],
      auxGrouping: Array[Int],
      namedProperties: Seq[NamedWindowProperty],
      mq: RelMetadataQuery,
      columns: ImmutableBitSet): ImmutableBitSet = {
    val fieldCount = window.getRowType.getFieldCount
    val columnList = columns.toList
    val groupingInToOutMap = new mutable.HashMap[Integer, Integer]()
    columnList.foreach { column =>
      require(column < fieldCount)
      if (column < grouping.length) {
        groupingInToOutMap.put(grouping(column), column)
      }
    }
    if (groupingInToOutMap.isEmpty) {
      columns
    } else {
      val fmq = FlinkRelMetadataQuery.reuseOrCreate(mq)
      val inputColumns = ImmutableBitSet.of(groupingInToOutMap.keys.toList)
      val inputUniqueGroups = fmq.getUniqueGroups(window.getInput, inputColumns)
      val uniqueGroupsFromGrouping = inputUniqueGroups.asList.map { i =>
        groupingInToOutMap.getOrElse(i, throw new IllegalArgumentException(s"Illegal index: $i"))
      }
      if (columns.equals(ImmutableBitSet.of(grouping ++ auxGrouping: _*))) {
        return ImmutableBitSet.of(uniqueGroupsFromGrouping)
      }

      val groupingOutCols = groupingInToOutMap.values
      // TODO drop some nonGroupingCols base on FlinkRelMdColumnUniqueness#areColumnsUnique(window)
      val nonGroupingCols = columnList.filterNot(groupingOutCols.contains)
      ImmutableBitSet.of(uniqueGroupsFromGrouping).union(ImmutableBitSet.of(nonGroupingCols))
    }
  }

  def getUniqueGroups(
      over: Window,
      mq: RelMetadataQuery,
      columns: ImmutableBitSet): ImmutableBitSet = {
    getUniqueGroupsOfOver(over.getRowType.getFieldCount, over.getInput, mq, columns)
  }

  def getUniqueGroups(
      over: BatchExecOverAggregate,
      mq: RelMetadataQuery,
      columns: ImmutableBitSet): ImmutableBitSet = {
    getUniqueGroupsOfOver(over.getRowType.getFieldCount, over.getInput, mq, columns)
  }

  private def getUniqueGroupsOfOver(
      outputFiledCount: Int,
      input: RelNode,
      mq: RelMetadataQuery,
      columns: ImmutableBitSet): ImmutableBitSet = {
    val inputFieldCount = input.getRowType.getFieldCount
    val (inputColumns, nonInputColumns) = columns.toList.partition(_ < inputFieldCount)
    val fmq = FlinkRelMetadataQuery.reuseOrCreate(mq)
    val inputUniqueGroups = fmq.getUniqueGroups(input, ImmutableBitSet.of(inputColumns))
    inputUniqueGroups.union(ImmutableBitSet.of(nonInputColumns))
  }

  // Catch-all rule when none of the others apply.
  def getUniqueGroups(
      rel: RelNode,
      mq: RelMetadataQuery,
      columns: ImmutableBitSet): ImmutableBitSet = columns

  def getUniqueGroups(
      rel: RelSubset,
      mq: RelMetadataQuery,
      columns: ImmutableBitSet): ImmutableBitSet = {
    if (!Bug.CALCITE_1048_FIXED) {
      //if the best node is null, so we can get the uniqueKeys based original node, due to
      //the original node is logically equivalent as the rel.
      val fmq = FlinkRelMetadataQuery.reuseOrCreate(mq)
      fmq.getUniqueGroups(Util.first(rel.getBest, rel.getOriginal), columns)
    } else {
      throw new RuntimeException("CALCITE_1048 is fixed, so check this method again!")
    }
  }

}

object FlinkRelMdUniqueGroups {

  private val INSTANCE = new FlinkRelMdUniqueGroups

  val SOURCE: RelMetadataProvider = ReflectiveRelMetadataProvider.reflectiveSource(
    FlinkMetadata.UniqueGroups.METHOD, INSTANCE)

}