package tech.mlsql.job

abstract class JobListener {

  import JobListener._

  def onJobStarted(event: JobStartedEvent): Unit

  def onJobFinished(event: JobFinishedEvent): Unit

}

object JobListener {

  trait JobEvent

  class JobStartedEvent(val groupId:String) extends JobEvent

  class JobFinishedEvent(val groupId:String) extends JobEvent

}

object SQLCacheExt {
  val cache = new java.util.concurrent.ConcurrentHashMap[String, ArrayBuffer[TableCacheItem]]()

  def addCache(tci: TableCacheItem) = {
    synchronized {
      val items = cache.getOrDefault(tci.groupId, ArrayBuffer[TableCacheItem]())
      items += tci
      cache.put(tci.groupId, items)
    }
  }

  def cleanCache(session: SparkSession, groupId: String) = {
    val items = cache.remove(groupId)
    if (items != null) {
      items.foreach { item =>
        SparkExposure.cleanCache(session, item.planToCache)
      }
    }
  }
}

case class TableCacheItem(groupId: String, tableName: String, planToCache: LogicalPlan, cacheStartTime: Long)

val df = if (path.isEmpty) _df else _df.sparkSession.table(path)
    val exe = params.get(execute.name).getOrElse {
      "cache"
    }
    val __dfname__ = params("__dfname__")
    val _isEager = params.get(isEager.name).map(f => f.toBoolean).getOrElse(false)

    if (!execute.isValid(exe)) {
      throw new MLSQLException(s"${execute.name} should be cache or uncache")
    }

    val context = ScriptSQLExec.contextGetOrForTest()

    if (exe == "cache") {
      df.persist()
      SQLCacheExt.addCache(TableCacheItem(context.groupId, context.owner, __dfname__, df.queryExecution.logical, System.currentTimeMillis()))
    } else {
      df.unpersist()
    }

    if (_isEager) {
      df.count()
    }
    df

def run(session: SparkSession, job: MLSQLJobInfo, f: () => Unit): Unit = {
    try {
      _jobListeners.foreach { f => f.onJobStarted(new JobStartedEvent(job.groupId)) }
      if (_jobManager == null) {
        f()
      } else {
        session.sparkContext.setJobGroup(job.groupId, job.jobName, true)
        _jobManager.groupIdToMLSQLJobInfo.put(job.groupId, job)
        f()
      }
      _jobListeners.foreach { f => f.onJobFinished(new JobFinishedEvent(job.groupId)) }
    } finally {
      handleJobDone(job.groupId)
      session.sparkContext.clearJobGroup()
    }
  }