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

import org.apache.datasketches.hll.HllSketch
import org.apache.datasketches.quantiles.DoublesSketch
import org.apache.datasketches.tuple.arrayofdoubles.ArrayOfDoublesSketch
import org.apache.druid.query.aggregation.datasketches.hll.HllSketchModule
import org.apache.druid.query.aggregation.datasketches.quantiles.DoublesSketchModule
import org.apache.druid.query.aggregation.datasketches.theta.{SketchHolder, SketchModule}
import org.apache.druid.query.aggregation.datasketches.tuple.ArrayOfDoublesSketchModule
import org.apache.druid.query.aggregation.histogram.{ApproximateHistogram,
  ApproximateHistogramDruidModule, FixedBucketsHistogram, FixedBucketsHistogramAggregator}
import org.apache.druid.query.aggregation.variance.{VarianceAggregatorCollector, VarianceSerde}
import org.apache.druid.segment.serde.ComplexMetrics
import org.apache.druid.spark.mixins.Logging

import scala.collection.mutable

/**
 * A registry for plugging in support for Druid complex types. Provides definitions for supporting complex types in
  * extensions-core out of the box.
 */
object ComplexTypeRegistry extends Logging {
  private val registeredSerdeInitFunctions: mutable.HashMap[String, () => Unit] = new mutable.HashMap()
  private val registeredSerializeFunctions: mutable.HashMap[Class[_], Any => Array[Byte]] =
    new mutable.HashMap()

  /**
    * Register a function REGISTERSERDEFUNC that initializes serializers and deserializers for the complex type with
    * the name NAME. Assumes that the associated complex type is serialized as a byte array.
    *
    * @param name The type name of the complex type to register.
    * @param registerSerdeFunc The function to use to register the necessary serdes for this complex type.
    */
  def register(
                name: String,
                registerSerdeFunc: () => Unit
              ): Unit = {
    logInfo(s"Registering serde initializers for complex type $name.")
    registeredSerdeInitFunctions(name) = registerSerdeFunc
  }

  /**
    * Register a function REGISTERSERDEFUNC that initializes serializers and deserializers, a class DESERIALIZEDCLASS
    * to deserialize byte arrays into, and a function SERIALIZEFUNC for converting from instances of DESERIALIZEDCLASS
    * to byte arrays for the complex type with the name NAME. Assumes that the associated complex type is
    * serialized as a byte array.
    *
    * @param name The type name of the complex type to register.
    * @param registerSerdeFunc The function to use to register the necessary serdes for this complex type.
    * @param deserializedClass The class to deserialize values of the registered type to from byte arrays.
    * @param serializeFunc The function to use when serializing instances of DESERIALIZEDCLASS to byte arrays.
    */
  def register(
                name: String,
                registerSerdeFunc: () => Unit,
                deserializedClass: Class[_],
                serializeFunc: Any => Array[Byte]): Unit = {
    logInfo(s"Registering serde initializers and serialization functions for complex type $name.")
    registeredSerdeInitFunctions(name) = registerSerdeFunc
    registeredSerializeFunctions(deserializedClass) = serializeFunc
  }

  /**
    * Shortcut for registering known complex type serdes (e.g. those in extensions-core) by name.
    *
    * @param name The type name of the complex type to register.
    * @param shouldCompact Whether or not to store compacted versions of this complex type. Ignored for complex types
    *                      that don't have compacted forms.
    */
  def registerByName(name: String, shouldCompact: Boolean = false): Unit = {
    if (!registeredSerdeInitFunctions.contains(name) && knownTypes.contains(name)) {
      knownTypes(name)(shouldCompact)
    }
  }

  def getRegisteredMetricNames: Set[String] = {
    registeredSerdeInitFunctions.keySet.toSet
  }

  def getRegisteredSerializedClasses: Set[Class[_]] = {
    registeredSerializeFunctions.keySet.toSet
  }

  def registerSerdeInitFunctions(complexTypeName: String): Unit = {
    if (registeredSerdeInitFunctions.contains(complexTypeName)) {
      registeredSerdeInitFunctions(complexTypeName).apply()
    }
  }

  def deserialize(col: Any): Array[Byte] = {
    if (registeredSerializeFunctions.keySet.contains(col.getClass)) {
      registeredSerializeFunctions(col.getClass)(col)
    } else {
      throw new IllegalArgumentException(
        s"Unsure how to parse ${col.getClass.toString} into a ByteArray!"
      )
    }
  }

  def registerSerdes(): Unit = {
    registeredSerdeInitFunctions.foreach(_._2.apply())
  }

  /**
    * Register serdes for all complex types in extensions-core.
    */
  def initializeDefaults(shouldCompact: Boolean = false): Unit = {
    knownTypes.foreach(_._2.apply(shouldCompact))
  }

  private val knownTypes: Map[String, Boolean => Unit] = Map[String, Boolean => Unit](
    // Approximate Histograms
    "approximateHistogram" -> ((_: Boolean) =>
      register(
        "approximateHistogram",
        () => ApproximateHistogramDruidModule.registerSerde(),
        classOf[ApproximateHistogram],
        histogram => histogram.asInstanceOf[ApproximateHistogram].toBytes
      )),
    // Fixed Bucket Histograms
    FixedBucketsHistogramAggregator.TYPE_NAME -> ((_: Boolean) =>
        register(
        FixedBucketsHistogramAggregator.TYPE_NAME,
        () => ApproximateHistogramDruidModule.registerSerde(),
        classOf[FixedBucketsHistogram],
        histogram => histogram.asInstanceOf[FixedBucketsHistogram].toBytes
      )),
    // Tuple Sketches
    ArrayOfDoublesSketchModule.ARRAY_OF_DOUBLES_SKETCH -> ((_: Boolean) =>
      register(
          ArrayOfDoublesSketchModule.ARRAY_OF_DOUBLES_SKETCH,
          // TODO: This probably needs to be wrapped in a try to ensure it only happens once
          () => new ArrayOfDoublesSketchModule().configure(null), // scalastyle:ignore null
          classOf[ArrayOfDoublesSketch],
          sketch => sketch.asInstanceOf[ArrayOfDoublesSketch].toByteArray
      )),
    // Quantiles Sketches
    DoublesSketchModule.DOUBLES_SKETCH -> ((shouldCompact: Boolean) =>
      register(
        DoublesSketchModule.DOUBLES_SKETCH,
        () => DoublesSketchModule.registerSerde(),
        classOf[DoublesSketch],
        sketch => {
          val doublesSketch = sketch.asInstanceOf[DoublesSketch]
          if (shouldCompact) doublesSketch.toByteArray(shouldCompact) else doublesSketch.toByteArray
        }
      )),
    // HLL Sketches
    HllSketchModule.TYPE_NAME -> ((shouldCompact: Boolean) =>
      register(
        HllSketchModule.TYPE_NAME,
        () => HllSketchModule.registerSerde(),
        classOf[HllSketch],
        sketch => {
          val hllSketch = sketch.asInstanceOf[HllSketch]
          if (shouldCompact) hllSketch.toCompactByteArray else hllSketch.toUpdatableByteArray
        }
      )),
    // Theta Sketches
    SketchModule.THETA_SKETCH -> ((shouldCompact: Boolean) =>
      register(
        SketchModule.THETA_SKETCH,
        () => SketchModule.registerSerde(),
        classOf[SketchHolder],
        sketch => {
          val thetaSketch = sketch.asInstanceOf[SketchHolder].getSketch
          if (shouldCompact) thetaSketch.compact().toByteArray else thetaSketch.toByteArray
        }
      )),
    SketchModule.THETA_SKETCH_BUILD_AGG -> ((shouldCompact: Boolean) =>
      register(
        SketchModule.THETA_SKETCH_BUILD_AGG,
        () => SketchModule.registerSerde(),
        classOf[SketchHolder], // TODO: Maybe?
        sketch => {
          val thetaSketch = sketch.asInstanceOf[SketchHolder].getSketch
          if (shouldCompact) thetaSketch.compact().toByteArray else thetaSketch.toByteArray
        }
      )),
    SketchModule.THETA_SKETCH_MERGE_AGG -> ((shouldCompact: Boolean) =>
      register(
        SketchModule.THETA_SKETCH_MERGE_AGG,
        () => SketchModule.registerSerde(),
        classOf[SketchHolder], // TODO: Maybe?
        sketch => {
          val thetaSketch = sketch.asInstanceOf[SketchHolder].getSketch
          if (shouldCompact) thetaSketch.compact().toByteArray else thetaSketch.toByteArray
        }
      )),
    // Variance
    "variance" -> ((_: Boolean) =>
      register(
        "variance",
        () => ComplexMetrics.registerSerde("variance", new VarianceSerde()),
        classOf[VarianceAggregatorCollector],
        collector => collector.asInstanceOf[VarianceAggregatorCollector].toByteArray
      ))
  )
}