package com.lightbend.scala.kafkastreams.queriablestate.withstore

import javax.ws.rs.NotFoundException
import akka.actor.ActorSystem
import akka.http.scaladsl.Http
import akka.http.scaladsl.server.Directives._
import akka.http.scaladsl.server.Route
import akka.stream.ActorMaterializer
import akka.util.Timeout
import com.lightbend.scala.kafkastreams.queriablestate.MetadataService
import com.lightbend.java.configuration.kafka.ApplicationKafkaParameters._
import com.lightbend.scala.kafkastreams.store.StoreState
import com.lightbend.scala.kafkastreams.store.store.custom.ModelStateStoreType
import de.heikoseeberger.akkahttpjackson.JacksonSupport
import org.apache.kafka.streams.KafkaStreams
import org.apache.kafka.streams.state.QueryableStoreTypes

import scala.concurrent.duration._

/**
 *  A simple REST proxy that runs embedded in the Model server. This is used to
 *  demonstrate how a developer can use the Interactive Queries APIs exposed by Kafka Streams to
 *  locate and query the State Stores within a Kafka Streams Application.
 *  @see https://github.com/confluentinc/examples/blob/3.2.x/kafka-streams/src/main/java/io/confluent/examples/streams/interactivequeries/WordCountInteractiveQueriesRestService.java
 */
object RestServiceStore {

  implicit val system = ActorSystem("ModelServing")
  implicit val materializer = ActorMaterializer()
  implicit val executionContext = system.dispatcher
  implicit val timeout = Timeout(10.seconds)
  val host = "127.0.0.1"
  val port = 8888

  def startRestProxy(streams: KafkaStreams, port: Int, storeType : String) = {

    val routes: Route = QueriesResource.storeRoutes(streams, port, storeType)
    Http().bindAndHandle(routes, host, port) map
      { binding => println(s"Starting models observer on port ${binding.localAddress}") } recover {
      case ex =>
        println(s"Models observer could not bind to $host:$port - ${ex.getMessage}")
    }
  }
}

object QueriesResource extends JacksonSupport {

  private val customStoreType = new ModelStateStoreType()
  private val standardStoreType = QueryableStoreTypes.keyValueStore[Integer,StoreState]

  def storeRoutes(streams: KafkaStreams, port : Int, storeType : String): Route = {
    val metadataService = new MetadataService(streams)
    get {
      pathPrefix("state") {
        path("instances") {
          complete(
            metadataService.streamsMetadataForStore(STORE_NAME, port)
          )
        } ~
          path("value") {
            storeType match {
              case "custom" =>
                val store = streams.store(STORE_NAME, customStoreType)
                if (store == null) throw new NotFoundException
                complete(store.getCurrentServingInfo)
              case _ =>
                val store = streams.store(STORE_NAME, standardStoreType)
                if (store == null) throw new NotFoundException
                complete(store.get(STORE_ID).currentState)
            }
          }
        }
    }
  }
}