package kotlinx.coroutines

import java.util.*

class RootIterator<T>(iterator: NestedIterator<T>) : AbstractIterator<T>() {
    private val stack = Stack<Iterator<T>>().apply { push(iterator) }

    override fun computeNext() {
        while (true) {
            if (stack.isEmpty()) {
                done()
                return
            }
            val i = stack.peek()
            if (!i.hasNext()) {
                stack.pop()
            } else {
                if (i is NestedIterator<T> && i.nextNestedIterator != null) {
                    stack.push(i.nextNestedIterator)
                    i.next()
                } else {
                    setNext(i.next())
                    return
                }
            }
        }
    }
}

interface NestedIterable<T> : Sequence<T> {
    fun nestedIterator(): NestedIterator<T>
    override fun iterator(): Iterator<T> = RootIterator(nestedIterator())
}

abstract class NestedIterator<T> : AbstractIterator<T>() {
    var nextNestedIterator: Iterator<T>? = null
        private set

    fun providesIterator() = hasNext() && nextNestedIterator != null

    final override fun computeNext(): Unit {
        nextNestedIterator = null
        computeNextItemOrIterator()
    }

    abstract fun computeNextItemOrIterator()

    protected fun setNextIterator(iterator: Iterator<T>) {
        nextNestedIterator = iterator
        setNext(null as T) //state transfer to Ready
    }
}

/**
 * Creates a Sequence object based on received coroutine [c].
 *
 * Each call of 'yield' suspend function within the coroutine lambda generates
 * next element of resulting sequence.
 */
fun <T> generate(coroutine c: GeneratorController<T>.() -> Continuation<Unit>
) = object : NestedIterable<T> {
    override fun nestedIterator(): NestedIterator<T> {
        val iterator = GeneratorController<T>()
        iterator.setNextStep(iterator.c())
        return iterator
    }
}

class GeneratorController<T> internal constructor() : NestedIterator<T>() {
    private lateinit var nextStep: Continuation<Unit>

    override fun computeNextItemOrIterator() {
        nextStep.resume(Unit)
    }

    internal fun setNextStep(step: Continuation<Unit>) {
        nextStep = step
    }

    suspend fun yield(value: T, c: Continuation<Unit>) {
        setNext(value)
        setNextStep(c)
    }

    suspend fun yieldAll(values: Sequence<T>, c: Continuation<Unit>) {
        setNextIterator(if (values is NestedIterable<T>)
                            values.nestedIterator() else
                            values.iterator())
        setNextStep(c)
    }

    operator fun handleResult(result: Unit, c: Continuation<Nothing>) {
        done()
    }
}