Maintenance.scala

import cats.effect.IO
import cats.effect.unsafe.implicits.global

import scala.concurrent.duration._

object Maintenance {
  def schedule(schedule: FiniteDuration,
               initialDelay: Option[FiniteDuration] = None)
              (action: IO[TaskResult]): MaintenanceTask = {
    var normalSchedule = schedule
    var stat = TaskStatus().schedule(initialDelay.getOrElse(schedule))
    var cancelled = false
    val task = new MaintenanceTask {
      override def status: TaskStatus = stat

      override def cancel(): Unit = cancelled = true
    }

    def scheduleNext(resultOption: Option[TaskResult]): IO[TaskResult] = {
      stat = stat.copy(lastRun = Some(System.currentTimeMillis()), timesRun = stat.timesRun + 1)
      if (cancelled) {
        stat = stat.copy(nextRun = None, nextSchedule = None)
        IO.pure(TaskResult.Stop)
      } else {
        val nextRunOption = resultOption match {
          case None => Some(initialDelay.getOrElse(schedule))
          case Some(TaskResult.Continue) => Some(normalSchedule)
          case Some(TaskResult.Stop) => None
          case Some(TaskResult.RunAgain) => Some(0.seconds)
          case Some(TaskResult.ChangeSchedule(delay)) =>
            normalSchedule = delay
            Some(delay)
          case Some(TaskResult.NextSchedule(delay)) => Some(delay)
        }
        nextRunOption match {
          case Some(nextRun) =>
            stat = stat.schedule(nextRun)
            action.flatMap { result =>
              scheduleNext(Some(result))
            }
          case None =>
            stat = stat.copy(nextRun = None, nextSchedule = None)
            IO.pure(TaskResult.Stop)
        }
      }
    }

    scheduleNext(None).unsafeRunAndForget()

    task
  }
}

trait MaintenanceTask {
  def status: TaskStatus

  def cancel(): Unit
}

case class TaskStatus(lastRun: Option[Long] = None,
                      timesRun: Int = 0,
                      nextRun: Option[Long] = None,
                      nextSchedule: Option[FiniteDuration] = None) {
  def schedule(duration: FiniteDuration): TaskStatus = copy(
    nextRun = Some(System.currentTimeMillis() + duration.toMillis),
    nextSchedule = Some(duration)
  )
}

sealed trait TaskResult

object TaskResult {
  case object Continue extends TaskResult
  case object Stop extends TaskResult
  case object RunAgain extends TaskResult
  case class ChangeSchedule(delay: FiniteDuration) extends TaskResult
  case class NextSchedule(delay: FiniteDuration) extends TaskResult
}