ElasticDslExt.scala

import java.time.ZoneOffset

import com.sksamuel.elastic4s.admin.CreateIndexTemplateDefinition
import com.sksamuel.elastic4s.mappings.{DynamicTemplateDefinition, TypedFieldDefinition}
import com.sksamuel.elastic4s.searches.QueryDefinition
import com.sksamuel.elastic4s.searches.aggs.{AggregationDefinition, FiltersAggregationDefinition, RangeAggregationDefinition}
import com.sksamuel.elastic4s.{ElasticDsl, IndexAndType, Indexable}
import org.elasticsearch.action.DocWriteResponse.Result
import org.elasticsearch.action.update.UpdateResponse
import org.elasticsearch.search.aggregations.bucket.histogram.{DateHistogramInterval, ExtendedBounds}
import org.joda.time.DateTimeZone

import scala.concurrent.duration._
import scala.language.implicitConversions

object ElasticDslExt extends ElasticDsl {

  implicit object StringIndexable extends Indexable[String] {
    override def json(js: String): String = js
  }

  implicit class RangeAggregationDefinitionExt(val aggr: RangeAggregationDefinition) extends AnyVal {
    def field(field: String) = {
      aggr.builder.field(field)
      aggr
    }
  }

  implicit def stringTupleToIndexAndType(indexAndType: (String, String)): IndexAndType = IndexAndType(indexAndType._1, indexAndType._2)
  implicit def stringTupleToExtendedBounds(fromTo: (String, String)): ExtendedBounds = new ExtendedBounds(fromTo._1, fromTo._2)
  implicit def dateHistogramIntervalToMillis(interval: DateHistogramInterval): Long = interval match {
    case DateHistogramInterval.DAY    ⇒ 1.day.toMillis
    case DateHistogramInterval.HOUR   ⇒ 1.hour.toMillis
    case DateHistogramInterval.MINUTE ⇒ 1.minute.toMillis
    case DateHistogramInterval.SECOND ⇒ 1.second.toMillis
    case DateHistogramInterval.WEEK   ⇒ 7.days.toMillis
    case x                            ⇒ throw new NotImplementedError(s"don't know how to convert to millis: $x")
  }

  implicit def javaToJodaTimeZone(tz: ZoneOffset): DateTimeZone = DateTimeZone.forOffsetMillis(tz.getTotalSeconds * 1000)

  override def createTemplate(name: String) = new CreateIndexTemplateDefinitionBuilder(name)
  override def dynamicTemplate(nameOfTemplate: String) = new DynamicTemplateDefinitionBuilder(nameOfTemplate)

  class CreateIndexTemplateDefinitionBuilder(name: String) {
    def pattern(pat: String) = CreateIndexTemplateDefinition(name, pat)
  }

  class DynamicTemplateDefinitionBuilder(nameOfTemplate: String) {
    def mapping(mapping: TypedFieldDefinition) = DynamicTemplateDefinition(nameOfTemplate, mapping)
  }

  implicit class UpdateResponseExt(val r: UpdateResponse) extends AnyVal {
    def created: Boolean = r.getResult == Result.CREATED
  }

  override def filtersAggregation(name: String) = new FiltersAggregationDefinitionExt(name)

  class FiltersAggregationDefinitionExt(name: String) {
    def queries(first: QueryDefinition, rest: QueryDefinition*): FiltersAggregationDefinition = queries(first +: rest)
    def queries(queries: Iterable[QueryDefinition]): FiltersAggregationDefinition =
      FiltersAggregationDefinition(name, queries)
  }

  implicit class AggregationDefinitionExt[D <: AggregationDefinition](val d: D) extends AnyVal {
    def aggregations(it: Iterable[AggregationDefinition]): D = {
      it.foreach { aad ⇒ d.builder.subAggregation(aad.builder) }
      d
    }

    def aggregations(a: AggregationDefinition*): D = aggregations(a.toIterable)

    def aggs(a: AggregationDefinition*): D = aggregations(a)

    def aggs(a: Iterable[AggregationDefinition]): D = aggregations(a)
  }

}