import kotlinx.coroutines.SupervisorJob
import kotlinx.coroutines.cancel
import processing.core.PApplet
import java.util.SortedSet
import java.util.TreeSet
import kotlin.coroutines.Continuation
import kotlin.coroutines.createCoroutine
import kotlin.coroutines.intrinsics.COROUTINE_SUSPENDED
import kotlin.coroutines.intrinsics.suspendCoroutineUninterceptedOrReturn
import kotlin.coroutines.resume

data class ScheduledCallback(
    val suspendedAtMillis: Int,
    val resumeAtMillis: Int,
    val continuation: Continuation<Unit>
) : Comparable<ScheduledCallback> {
    override fun compareTo(other: ScheduledCallback): Int {
        val resumeAtComparison = resumeAtMillis.compareTo(other.resumeAtMillis)

        return if (resumeAtComparison == 0)
            suspendedAtMillis.compareTo(other.suspendedAtMillis)
        else
            resumeAtComparison
    }
}

open class ScheduledSketch : PApplet() {
    /**
     * A set of pending tasks. Each task has an associated millisecond value where it will be resumed as soon as
     * possible. Every time draw() is called, every callback in this set whose millisecond value is less than or equal
     * to the current millisecond value will be executed and removed from the set.
     */
    private val schedule: SortedSet<ScheduledCallback> = TreeSet()

    /**
     * The job that all running tasks will be grouped under, so that once the application is closed, everything is
     * cleaned up.
     */
    private val job = SupervisorJob()

    override fun exitActual() {
        job.cancel("The sketch is exiting!")
        super.exitActual()
    }

    override fun setup() {
        this::main.createCoroutine(Continuation(job, ::println)).resume(Unit)
    }

    suspend fun main() {
        while (true) {
            background(255f, 255f, 255f)
            circle(128f, 128f, 16f)
            yieldFor(20)
            background(255f, 255f, 255f)
            yieldFor(20)
        }
    }

    private suspend fun yieldUntil(millis: Int) = suspendCoroutineUninterceptedOrReturn<Unit> {
        schedule.add(ScheduledCallback(millis(), millis, it))
        COROUTINE_SUSPENDED
    }

    private suspend fun yieldFor(millis: Int) = yieldUntil(millis() + millis)

    override fun draw() {
        val millis = millis()
        val queue = TreeSet<ScheduledCallback>()

        schedule.forEach { scheduled ->
            if (millis >= scheduled.resumeAtMillis) {
                queue.add(scheduled)
                schedule.remove(scheduled)
            }
        }

        queue.forEach { scheduled ->
            scheduled.continuation.resume(Unit)
        }
    }
}