import androidx.compose.runtime.*
import androidx.compose.runtime.snapshots.Snapshot
import androidx.compose.ui.test.TestMonotonicFrameClock
import kotlinx.coroutines.CoroutineScope
import kotlinx.coroutines.ExperimentalCoroutinesApi
import kotlinx.coroutines.coroutineScope
import kotlinx.coroutines.launch
import kotlinx.coroutines.test.runBlockingTest
import kotlin.coroutines.CoroutineContext
import kotlin.coroutines.EmptyCoroutineContext

private fun callSetContent(composition: Composition, content: @Composable () -> Unit) {
    composition.setContent(content)
}

@OptIn(InternalComposeApi::class)
@Composable
fun TestSubcomposition(
    content: @Composable () -> Unit
) {
    val parentRef = rememberCompositionContext()
    val currentContent by rememberUpdatedState(content)
    DisposableEffect(parentRef) {
        val subcomposition = Composition(EmptyApplier(), parentRef)
        // TODO: work around for b/179701728
        callSetContent(subcomposition) {
            // Note: This is in a lambda invocation to keep the currentContent state read
            // in the subcomposition's content composable. Changing this to be
            // subcomposition.setContent(currentContent) would snapshot read only on initial set.
            currentContent()
        }
        onDispose {
            subcomposition.dispose()
        }
    }
}

class EmptyApplier : Applier<Unit> {
    override val current: Unit = Unit
    override fun down(node: Unit) {}
    override fun up() {}
    override fun insertTopDown(index: Int, instance: Unit) {
        error("Unexpected")
    }
    override fun insertBottomUp(index: Int, instance: Unit) {
        error("Unexpected")
    }
    override fun remove(index: Int, count: Int) {
        error("Unexpected")
    }
    override fun move(from: Int, to: Int, count: Int) {
        error("Unexpected")
    }
    override fun clear() {}
}

@OptIn(ExperimentalCoroutinesApi::class)
suspend fun <R> localRecomposerTest(
    block: CoroutineScope.(Recomposer) -> R
) = coroutineScope {
    val contextWithClock = coroutineContext + TestMonotonicFrameClock(this)
    val recomposer = Recomposer(contextWithClock)
    launch(contextWithClock) {
        recomposer.runRecomposeAndApplyChanges()
    }
    block(recomposer)
    // This call doesn't need to be in a finally; everything it does will be torn down
    // in exceptional cases by the coroutineScope failure
    recomposer.cancel()
}

fun runBlockingCompositionTest(
    context: CoroutineContext = EmptyCoroutineContext,
    composable: @Composable () -> Unit
) {
    runBlockingTest(context = context) {
        localRecomposerTest { recomposer ->
            val composition = Composition(EmptyApplier(), recomposer)

            Snapshot.notifyObjectsInitialized()
            composition.setContent {
                TestSubcomposition {
                    composable()
                }
            }

            Snapshot.sendApplyNotifications()
            advanceUntilIdle()
        }
    }
}