/*
 * Copyright 2016 Tuplejump
 *
 * Licensed 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 com.tuplejump.embedded.kafka

import java.io.{File => JFile}
import java.net.InetSocketAddress
import java.util.concurrent.atomic.AtomicReference

import scala.util.Try
import org.I0Itec.zkclient.exception.ZkMarshallingError
import org.I0Itec.zkclient.serialize.ZkSerializer
import org.apache.zookeeper.server.{ NIOServerCnxnFactory, ZooKeeperServer }

/**
 * Implements a simple standalone ZooKeeperServer.
 * To create a ZooKeeper client object, the application needs to pass a
 * connection string containing a comma separated list of host:port pairs,
 * each corresponding to a ZooKeeper server.
 * <p>
 * Session establishment is asynchronous. This constructor will initiate
 * connection to the server and return immediately - potentially (usually)
 * before the session is fully established. The watcher argument specifies
 * the watcher that will be notified of any changes in state. This
 * notification can come at any point before or after the constructor call
 * has returned.
 * <p>
 * The instantiated ZooKeeper client object will pick an arbitrary server
 * from the connectString and attempt to connect to it. If establishment of
 * the connection fails, another server in the connect string will be tried
 * (the order is non-deterministic, as we random shuffle the list), until a
 * connection is established. The client will continue attempts until the
 * session is explicitly closed.
 */
class EmbeddedZookeeper(val connectTo: String,
                        val tickTime: Int,
                        snapDir: JFile,
                        dataDir: JFile) extends Assertions with Logging {

  logger.info("Starting Zookeeper")

  private val _factory = new AtomicReference[Option[NIOServerCnxnFactory]](None)

  private val _zookeeper = new AtomicReference[Option[ZooKeeperServer]](None)

  def isRunning: Boolean = _zookeeper.get exists (_.isRunning)

  def zookeeper: ZooKeeperServer = _zookeeper.get.getOrElse {
    logger.info("Attempt to call server before starting EmbeddedKafka instance. Starting automatically...")
    start()
    eventually(5000, 500)(assert(isRunning))

    _zookeeper.get.getOrElse(throw new IllegalStateException("Kafka server not initialized."))
  }

  /** Starts only one. */
  def start(): Unit = {
    val server = new ZooKeeperServer(snapDir, dataDir, tickTime)
    _zookeeper.set(Some(server))

    val (ip, port) = {
      val splits = connectTo.split(":")
      (splits(0), splits(1).toInt)
    }

    val f = new NIOServerCnxnFactory()
    f.configure(new InetSocketAddress(ip, port), 16)
    f.startup(server)

    _factory.set(Some(f))

    logger.info(s"ZooKeeperServer isRunning: $isRunning")
  }

  def shutdown(): Unit = {
    logger.info(s"Shutting down ZK NIOServerCnxnFactory.")

    for (v <- _factory.get) v.shutdown()
    _factory.set(None)

    for (v <- _zookeeper.get) {
      Try(v.shutdown())
      //awaitCond(!v.isRunning, 2000.millis)
      logger.info(s"ZooKeeper server shut down.")
    }
    _zookeeper.set(None)
  }
}

object DefaultStringSerializer extends ZkSerializer {

  @throws(classOf[ZkMarshallingError])
  def serialize(data: Object): Array[Byte] = data match {
    case a: String => a.getBytes("UTF-8")
    case _         => throw new ZkMarshallingError(s"Unsupported type '${data.getClass}'")
  }

  @throws(classOf[ZkMarshallingError])
  def deserialize(bytes: Array[Byte]): Object = bytes match {
    case b if Option(b).isEmpty => "" //ick
    case b                      => new String(bytes, "UTF-8")
  }
}