package ru.yandex.atom.utils.actor

import scala.collection.immutable.TreeSet
import scala.concurrent.duration.{FiniteDuration, Deadline}

/**
 * @author avhaliullin
 */
class Waiting[K, V] private(map: Map[K, V], queue: TreeSet[(K, Deadline)]) {
  def wait(k: K, v: V, duration: FiniteDuration): Waiting[K, V] =
    new Waiting[K, V](map + (k -> v), queue + (k -> duration.fromNow))

  def arrived(k: K): Option[(V, Waiting[K, V])] = {
    map.get(k).map {
      v =>
        (v, new Waiting(map - k, queue))
    }
  }

  def handleDead: (Waiting[K, V], Seq[(K, V)]) = {
    val (dead, rest) = queue.span(_._2.isOverdue())
    val rlyDead = dead.toSeq.map(_._1).filter(map.contains)

    (new Waiting[K, V](map -- rlyDead, rest), rlyDead.map(x => x -> map(x)))
  }

  def contains(k: K) = map.contains(k)

  def isEmpty = map.isEmpty

  def size = map.size
}

object Waiting {
  def apply[K, V](): Waiting[K, V] = {
    new Waiting[K, V](Map(), TreeSet[(K, Deadline)]()(Ordering.by(_._2)))
  }
}

