package ocelot.desktop.util

import scala.collection.mutable.ArrayBuffer

class Transaction protected () {
  private val cleanupCallbacks = ArrayBuffer.empty[() => Unit]

  def onFailure(f: => Unit): Unit = {
    cleanupCallbacks += f _
  }

  protected def runCleanupCallbacks(exc: Exception): exc.type = {
    for (callback <- cleanupCallbacks) {
      try {
        callback()
      } catch {
        case cleanupExc: Exception =>
          exc.addSuppressed(cleanupExc)
      }
    }

    exc
  }
}

object Transaction {
  private[Transaction] case class AbortedException(tx: AbortableTransaction) extends RuntimeException

  class AbortableTransaction protected[Transaction] () extends Transaction {
    def abort(): Nothing = {
      throw Transaction.AbortedException(this)
    }
  }

  def run[T](f: Transaction => T): T = {
    val tx = new Transaction

    try {
      f(tx)
    } catch {
      case exc: Exception => throw tx.runCleanupCallbacks(exc)
    }
  }

  def runAbortable[T](f: AbortableTransaction => T): Option[T] = {
    val tx = new AbortableTransaction

    try {
      Some(f(tx))
    } catch {
      case exc @ AbortedException(`tx`) =>
        val suppressed = tx.runCleanupCallbacks(exc).getSuppressed

        if (suppressed.nonEmpty) {
          throw exc
        }

        None
      case exc: Exception =>
        throw tx.runCleanupCallbacks(exc)
    }
  }
}
