package ru.yandex.atom.service

import akka.actor.{Props, Actor}
import ru.yandex.atom.zookeeper.ZookeeperActorComponent
import ru.yandex.atom.utils.collection._
import java.nio.ByteBuffer
import ru.yandex.atom.data.ReqID
import scala.concurrent.duration.FiniteDuration
import org.apache.zookeeper.{CreateMode, KeeperException}
import ru.yandex.atom.utils.actor.AtomActorLogging
import ru.yandex.atom.error.InternalProblem
import com.typesafe.config.Config

/**
 *
 * Эта реализация запрашивает id'шники через zookeeper пачками
 *
 * @author avhaliullin
 */
trait IdGeneratorActorComponentImpl extends IdGeneratorActorComponent {
  component: ZookeeperActorComponent =>

  class IdGeneratorActorImpl(config: IdGeneratorConfig) extends Actor with AtomActorLogging {

    import IdGeneratorRequest._
    import IdGeneratorResponse._

    private val ID_GENERATOR_VERSION: Byte = 1

    implicit val executionContext = context.dispatcher

    case object ReserveRequest

    case object PingReserve

    case class ReserveResponse(chunk: IdRange)

    val reserver = context.actorOf(Props(new ZookeeperRangeReservingActor), "zookeeper-range-reserver")

    var pool = IdPool()

    def receive = {
      case req@NextId(reqId) =>
        if (pool.isEmpty) {
          sender ! FailureResponse(req, InternalProblem.TemporarilyProblem)
          log.error("Failed to generate new id - the pool is empty")
        } else {
          val (newId, newPool) = pool.dequeue
          sender ! NextIdResponse(req, ID_GENERATOR_VERSION, newId)
          pool = newPool

          doReserve()
        }
      case PingReserve => doReserve()
      case ReserveResponse(chunk) => pool += chunk
    }

    def doReserve() {
      if (pool.size <= config.reserveWhenLeft) {
        reserver ! ReserveRequest
        context.system.scheduler.scheduleOnce(config.zookeeperTimeout * 2, self, PingReserve)
      }
    }


    override def preStart() = {
      super.preStart()
      doReserve()
    }

    class ZookeeperRangeReservingActor extends Actor with AtomActorLogging {
      private val zookeeperIdNode = "/url-list-id"

      case class PingRead(req: ZookeeperRequest.ReadRequest)

      case class PingUpdate(req: ZookeeperRequest.UpdateRequest)


      override def preStart() = {
        super.preStart()
        zookeeperActor ! ZookeeperRequest.CreateRequest(ReqID("ID range reserver", "pre-start"),
          zookeeperIdNode, toBytes(0L), CreateMode.PERSISTENT)
      }

      def receive = idle

      def idle: Receive = {
        case ReserveRequest => doRead()
      }

      def reading(readRequest: ZookeeperRequest.ReadRequest): Receive = {
        case ZookeeperResponse.ReadResponse(req, data, stat) =>
          val curId = fromBytes(data)
          val myLastId = curId + config.reserveSize
          val newRange = IdRange(curId + 1, config.reserveSize)
          log.info("Last id in zookeeper is {}, reserving range {}", curId, newRange)
          val updateRequest = ZookeeperRequest.UpdateRequest(newReqId, zookeeperIdNode, toBytes(myLastId), stat.version)
          doUpdate(updateRequest, newRange)

        case PingRead(r) if r == readRequest =>
          doRead(r)

        case ZookeeperResponse.FailureResponse(r, err) if r == readRequest =>
          log.error(err, "Zookeeper read failed")
          context.become(idle)
      }

      def updating(updateRequest: ZookeeperRequest, newRange: IdRange): Receive = {
        case ZookeeperResponse.UpdateResponse(req, stat) if req == updateRequest =>
          context.parent ! ReserveResponse(newRange)
          context.become(idle)

        case PingUpdate(r) if r == updateRequest =>
          doUpdate(r, newRange)

        case ZookeeperResponse.FailureResponse(r, err) if r == updateRequest =>
          err.code() match {
            case KeeperException.Code.BADVERSION =>
              log.warning("Zookeeper reported bad version")
              doRead()
            case _ =>
              log.error(err, "Zookeeper update failed")
              context.become(idle)
          }
      }

      def doRead(readRequest: ZookeeperRequest.ReadRequest = ZookeeperRequest.ReadRequest(newReqId, zookeeperIdNode)) = {
        zookeeperActor ! readRequest
        context.system.scheduler.scheduleOnce(config.zookeeperTimeout, self, PingRead(readRequest))
        context.become(reading(readRequest))
      }

      def doUpdate(updateRequest: ZookeeperRequest.UpdateRequest, newRange: IdRange) = {
        zookeeperActor ! updateRequest
        context.system.scheduler.scheduleOnce(config.zookeeperTimeout, self, PingUpdate(updateRequest))
        context.become(updating(updateRequest, newRange))
      }

      def fromBytes(data: Array[Byte]): Long = {
        ByteBuffer.wrap(data).getLong
      }

      def toBytes(id: Long): Array[Byte] = {
        val bb = ByteBuffer.allocate(8)
        bb.putLong(id)
        bb.array()
      }

      private var _reqId = 0L

      private def newReqId = {
        _reqId += 1
        ReqID(self.path.address.toString, _reqId.toString)
      }
    }

    case class IdRange(start: Long, size: Int) {
      def dequeue: (Long, IdRange) = (start, IdRange(start + 1, size - 1))

      def isEmpty = size == 0
    }

    case class IdPool(ranges: Vector[IdRange] = Vector()) {
      def hasNext: Boolean = !isEmpty

      lazy val size = ranges.foldLeft(0)(_ + _.size)

      def isEmpty = size == 0

      def +(range: IdRange) = IdPool(ranges :+ range)

      def +:(range: IdRange) = IdPool(range +: ranges)

      def dequeue: (Long, IdPool) = if (!hasNext) {
        throw new NoSuchElementException("Ids pool is empty")
      } else {
        val (range, rest) = ranges.dequeue
        val (id, rangeRest) = range.dequeue
        val newPool = if (rangeRest.isEmpty) rest else rangeRest +: rest
        id -> IdPool(newPool)
      }
    }


  }

  object IdGeneratorActorImpl {
    def props(config: IdGeneratorConfig) = Props(classOf[IdGeneratorActorImpl], component, config)
  }

  object IdGeneratorConfig {

    import ru.yandex.atom.utils.config._

    implicit def apply(config: Config): IdGeneratorConfig = IdGeneratorConfig(
      reserveSize = config.getInt("reserveSize"),
      reserveWhenLeft = config.getInt("reserveWhenLeft"),
      zookeeperTimeout = config.getFiniteDuration("zookeeperTimeout")
    )
  }

  case class IdGeneratorConfig(reserveSize: Int, reserveWhenLeft: Int, zookeeperTimeout: FiniteDuration)

}
