/*
 * 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.druid.spark.utils

import org.apache.druid.java.util.common.{ISE, JodaUtils}
import org.apache.druid.query.filter.{AndDimFilter, BoundDimFilter, DimFilter, FalseDimFilter,
  InDimFilter, NotDimFilter, OrDimFilter, RegexDimFilter, SelectorDimFilter}
import org.apache.druid.query.ordering.{StringComparator, StringComparators}
import org.apache.spark.sql.sources.{And, EqualNullSafe, EqualTo, Filter, GreaterThan,
  GreaterThanOrEqual, In, IsNotNull, IsNull, LessThan, LessThanOrEqual, Not, Or, StringContains,
  StringEndsWith, StringStartsWith}
import org.apache.spark.sql.types.{ArrayType, DataType, DoubleType, FloatType, IntegerType,
  LongType, StringType, StructType, TimestampType}

import scala.collection.JavaConverters.{seqAsJavaListConverter, setAsJavaSetConverter}

/**
  * Converters and utilities for working with Spark and Druid Filters.
  */
object FilterUtils {
  /**
    * Map an array of Spark filters FILTERS to a Druid dim filter or None if filters is empty.
    *
    * We return a DimFilter instead of a Filter and force callers to call .toFilter
    * or .toOptimizedFilter to get a filter because callers can't covert back to a DimFilter from a
    * Filter.
    *
    * @param filters The spark filters to map to a Druid filter.
    * @return A Druid filter corresponding to the union of filter conditions enumerated in FILTERS.
    */
  def mapFilters(filters: Array[Filter], schema: StructType): Option[DimFilter] = {
    if (filters.isEmpty) {
      Option.empty[DimFilter]
    } else {
      Some(new AndDimFilter(filters.map(mapFilter(_, schema)).toList.asJava).optimize())
    }
  }

  /**
    * Convert a Spark-style filter FILTER to a Druid-style filter.
    *
    * @param filter The Spark filter to map to a Druid filter.
    * @return The Druid filter corresponding to the filter condition described by FILTER.
    */
  def mapFilter(filter: Filter, schema: StructType): DimFilter = { // scalastyle:ignore method.length
    // scalastyle:off null
    filter match {
      case And(left, right) =>
        new AndDimFilter(List(mapFilter(left, schema), mapFilter(right, schema)).asJava)
      case Or(left, right) =>
        new OrDimFilter(List(mapFilter(left, schema), mapFilter(right, schema)).asJava)
      case Not(condition) =>
        new NotDimFilter(mapFilter(condition, schema))
      case IsNull(field) =>
        new SelectorDimFilter(field, null, null, null)
      case IsNotNull(field) => new NotDimFilter(new SelectorDimFilter(field, null, null, null))
      case In(field, values) =>
        new InDimFilter(field, values.filter(_ != null).map(_.toString).toSet.asJava, null, null)
      case StringContains(field, value) =>
        // Not 100% sure what Spark's expectations are for regex, case insensitive, etc.
        // and not sure the relative efficiency of various Druid dim filters
        // Could also use a SearchQueryDimFilter here
        // new LikeDimFilter(field, s"%$value%", null, null)
        new RegexDimFilter(field, value, null, null)
      case StringStartsWith(field, value) =>
        // Not sure the trade-offs between LikeDimFilter and RegexDimFilter here
        // new LikeDimFilter(field, s"$value%", null, null, null)
        new RegexDimFilter(field, s"^$value", null, null)
      case StringEndsWith(field, value) =>
        // Not sure the trade-offs between LikeDimFilter and RegexDimFilter here
        // new LikeDimFilter(field, s"%$value", null, null, null)
        new RegexDimFilter(field, s"$value$$", null, null)
      case EqualTo(field, value) =>
        if (value == null) {
          FalseDimFilter.instance()
        } else {
          new SelectorDimFilter(field, value.toString, null, null)
        }
      case EqualNullSafe(field, value) =>
        new SelectorDimFilter(field, Option(value).map(_.toString).orNull, null, null)
      case LessThan(field, value) =>
        new BoundDimFilter(
          field,
          null,
          value.toString,
          false,
          true,
          null,
          null,
          getOrderingFromDataType(schema(field).dataType),
          null
        )
      case LessThanOrEqual(field, value) =>
        new BoundDimFilter(
          field,
          null,
          value.toString,
          false,
          false,
          null,
          null,
          getOrderingFromDataType(schema(field).dataType),
          null
        )
      case GreaterThan(field, value) =>
        new BoundDimFilter(
          field,
          value.toString,
          null,
          true,
          false,
          null,
          null,
          getOrderingFromDataType(schema(field).dataType),
          null
        )
      case GreaterThanOrEqual(field, value) =>
        new BoundDimFilter(
          field,
          value.toString,
          null,
          false,
          false,
          null,
          null,
          getOrderingFromDataType(schema(field).dataType),
          null
        )
    }
    // scalastyle:on
  }

  private[utils] def getOrderingFromDataType(dataType: DataType): StringComparator = {
    dataType match {
      case LongType | IntegerType | DoubleType | FloatType => StringComparators.NUMERIC
      case StringType | ArrayType(StringType, _) => StringComparators.LEXICOGRAPHIC
      // Filters on complex types should return false when evaluated in isSupportedFilter, something's gone wrong
      case _ => throw new ISE("This reader doesn't support filtering on complex types! Complex type " +
        "filters should not be pushed down.")
    }
  }

  /**
    * Given a Spark filter and a target DataFrame schema, returns true iff the filter can be pushed down to the Druid
    * InputPartitionReaders. Since segments are pulled from deep storage before being filtered, this is not as useful
    * as it could be but still saves time and resources.
    *
    * @param filter The filter to evaluate for support.
    * @param schema The schema of the DataFrame to be filtered by FILTER.
    * @return True iff FILTER can be pushed down to the InputPartitionReaders.
    */
  private[spark] def isSupportedFilter(
                                        filter: Filter,
                                        schema: StructType,
                                        useSQLCompatibleNulls: Boolean = false
                                      ): Boolean = {
    // If the filter references columns we don't know about, we can't push it down
    if (!filter.references.forall(schema.fieldNames.contains(_))) {
      false
    } else {
      filter match {
        // scalastyle:off null
        case and: And => isSupportedFilter(and.left, schema, useSQLCompatibleNulls) &&
          isSupportedFilter(and.right, schema, useSQLCompatibleNulls)
        case or: Or => isSupportedFilter(or.left, schema, useSQLCompatibleNulls) &&
          isSupportedFilter(or.right, schema, useSQLCompatibleNulls)
        case not: Not => isSupportedFilter(not.child, schema, useSQLCompatibleNulls)
        // If we're using SQL-compatible nulls, we can filter for null.
        // Otherwise, callers should explictly filter for '' or 0 depending on the column type.
        // If we ever support pushing down filters on complex types, we'll need to add handling here.
        case _: IsNull => useSQLCompatibleNulls
        case _: IsNotNull => useSQLCompatibleNulls
        case in: In => checkAllDataTypesSupported(filter, schema) &&
          (useSQLCompatibleNulls || !in.values.contains(null))
        case _: StringContains => checkStringsOnly(filter, schema)
        case _: StringStartsWith => checkStringsOnly(filter, schema)
        case _: StringEndsWith => checkStringsOnly(filter, schema)
        // Hopefully Spark is smart enough to short-circuit for foo = NULL queries but if not, I guess we can
        case equalTo: EqualTo => checkAllDataTypesSupported(filter, schema) &&
          (useSQLCompatibleNulls || equalTo.value != null)
        case equalNullSafe: EqualNullSafe => checkAllDataTypesSupported(filter, schema) &&
          (useSQLCompatibleNulls || equalNullSafe.value != null)
        case _: LessThan => checkAllDataTypesSupported(filter, schema)
        case _: LessThanOrEqual => checkAllDataTypesSupported(filter, schema)
        case _: GreaterThan => checkAllDataTypesSupported(filter, schema)
        case _: GreaterThanOrEqual => checkAllDataTypesSupported(filter, schema)
        case _ => false
        // scalastyle:on
      }
    }
  }

  private def checkAllDataTypesSupported(filter: Filter, schema: StructType): Boolean = {
    filter.references.map{field =>
      schema(schema.fieldIndex(field)).dataType
    }.forall(supportedDataTypesForFiltering.contains(_))
  }

  private def checkStringsOnly(filter: Filter, schema: StructType): Boolean = {
    filter.references.map{field =>
      schema(schema.fieldIndex(field)).dataType == StringType
    }.forall(identity)
  }

  private val supportedDataTypesForFiltering: Seq[DataType] = Seq(
    IntegerType, LongType, FloatType, DoubleType, TimestampType, StringType
  )

  private val emptyBoundSeq = Seq.empty[(Bound, Long)]

  /**
    * Given an array of Spark filters, return upper and lower bounds on the value of the __time column if bounds
    * can be determined.
    *
    * @param filters The array of filters to extract __time bounds from
    * @return A tuple containing an optional lower and an optional upper bound on the __time column.
    */
  def getTimeFilterBounds(filters: Array[Filter]): (Option[Long], Option[Long]) = {
    val timeFilters = filters
      .filter(_.references.contains("__time"))
      .flatMap(FilterUtils.decomposeTimeFilters)
      .partition(_._1 == FilterUtils.LOWER)
    (timeFilters._1.map(_._2).reduceOption(_ max _),
      timeFilters._2.map(_._2).reduceOption(_ min _))
  }

  /**
    * Decompose a Spark Filter into a sequence of bounds on the __time field if possible.
    *
    * @param filter The Spark filter to possibly extract time bounds from.
    * @return A sequnce of tuples containing either UPPER or LOWER bounds on the __time field, in
    *         epoch millis.
    */
  private[spark] def decomposeTimeFilters(filter: Filter): Seq[(Bound, Long)] = { // scalastyle:ignore method.length
    filter match {
      case And(left, right) =>
        val bounds = Seq(left, right)
          .filter(_.references.contains("__time"))
          .flatMap(decomposeTimeFilters)
          .partition(_._1 == LOWER)
        val optBounds = (bounds._1.map(_._2).reduceOption(_ max _ ),
          bounds._2.map(_._2).reduceOption(_ min _ ))
        Seq[Option[(Bound, Long)]](
          optBounds._1.fold(Option.empty[(Bound, Long)])(bound => Some((LOWER, bound))),
          optBounds._2.fold(Option.empty[(Bound, Long)])(bound => Some((UPPER, bound)))
        ).flatten
      case Or(left, right) =>
        val bounds = Seq(left, right)
          .filter(_.references.contains("__time")).flatMap(decomposeTimeFilters)
          .partition(_._1 == LOWER)
        val optBounds = (bounds._1.map(_._2).reduceOption(_ min _ ),
          bounds._2.map(_._2).reduceOption(_ max _ ))
        Seq[Option[(Bound, Long)]](
          optBounds._1.fold(Option.empty[(Bound, Long)])(bound => Some((LOWER, bound))),
          optBounds._2.fold(Option.empty[(Bound, Long)])(bound => Some((UPPER, bound)))
        ).flatten
      case Not(condition) =>
        if (condition.references.contains("__time")) {
          // Our quick and dirty bounds enum doesn't handle nots, so just return an unbounded interval
          Seq[(Bound, Long)](
            (LOWER, JodaUtils.MIN_INSTANT),
            (UPPER, JodaUtils.MAX_INSTANT)
          )
        } else {
          emptyBoundSeq
        }
      case EqualTo(field, value) =>
        if (field == "__time") {
          Seq(
            (LOWER, value.asInstanceOf[Long]),
            (UPPER, value.asInstanceOf[Long])
          )
        } else {
          emptyBoundSeq
        }
      case LessThan(field, value) =>
        if (field == "__time") {
          Seq((UPPER, value.asInstanceOf[Long] - 1))
        } else {
          emptyBoundSeq
        }
      case LessThanOrEqual(field, value) =>
        if (field == "__time") {
          Seq((UPPER, value.asInstanceOf[Long]))
        } else {
          emptyBoundSeq
        }
      case GreaterThan(field, value) =>
        if (field == "__time") {
          Seq((LOWER, value.asInstanceOf[Long] + 1))
        } else {
          emptyBoundSeq
        }
      case GreaterThanOrEqual(field, value) =>
        if (field == "__time") {
          Seq((LOWER, value.asInstanceOf[Long]))
        } else {
          emptyBoundSeq
        }
      case _ => emptyBoundSeq
    }
  }

  private[spark] sealed trait Bound
  case object LOWER extends Bound
  case object UPPER extends Bound

}