package kotlinx.coroutines.flow

import kotlinx.coroutines.testing.*
import kotlinx.coroutines.*
import kotlin.coroutines.*
import kotlin.test.*
import kotlin.coroutines.coroutineContext as currentContext

class FlowContextOptimizationsTest : TestBase() {
    @Test
    fun testBaseline() = runTest {
        val flowDispatcher = wrapperDispatcher(currentContext)
        val collectContext = currentContext
        flow {
            assertSame(flowDispatcher, currentContext[ContinuationInterceptor] as CoroutineContext)
            expect(1)
            emit(1)
            expect(2)
            emit(2)
            expect(3)
        }
            .flowOn(flowDispatcher)
            .collect { value ->
                assertEquals(collectContext.minusKey(Job), currentContext.minusKey(Job))
                if (value == 1) expect(4)
                else expect(5)
            }

        finish(6)
    }

    @Test
    fun testFusedSameContext() = runTest {
        flow {
            expect(1)
            emit(1)
            expect(3)
            emit(2)
            expect(5)
        }
            .flowOn(currentContext.minusKey(Job))
            .collect { value ->
                if (value == 1) expect(2)
                else expect(4)
            }
        finish(6)
    }

    @Test
    fun testFusedSameContextWithIntermediateOperators() = runTest {
        flow {
            expect(1)
            emit(1)
            expect(3)
            emit(2)
            expect(5)
        }
            .flowOn(currentContext.minusKey(Job))
            .map { it }
            .flowOn(currentContext.minusKey(Job))
            .collect { value ->
                if (value == 1) expect(2)
                else expect(4)
            }
        finish(6)
    }

    @Test
    fun testFusedSameDispatcher() = runTest {
        flow {
            assertEquals("Name", currentContext[CoroutineName]?.name)
            expect(1)
            emit(1)
            expect(3)
            emit(2)
            expect(5)
        }
            .flowOn(CoroutineName("Name"))
            .collect { value ->
                assertNull(currentContext[CoroutineName]?.name)
                if (value == 1) expect(2)
                else expect(4)
            }
        finish(6)
    }

    @Test
    fun testFusedManySameDispatcher() = runTest {
        flow {
            assertEquals("Name1", currentContext[CoroutineName]?.name)
            assertEquals("OK", currentContext[CustomContextElement]?.str)
            expect(1)
            emit(1)
            expect(3)
            emit(2)
            expect(5)
        }
            .flowOn(CoroutineName("Name1")) // the first one works
            .flowOn(CoroutineName("Name2"))
            .flowOn(CoroutineName("Name3") + CustomContextElement("OK")) // but this is not lost
            .collect { value ->
                assertNull(currentContext[CoroutineName]?.name)
                assertNull(currentContext[CustomContextElement]?.str)
                if (value == 1) expect(2)
                else expect(4)
            }
        finish(6)
    }

    data class CustomContextElement(val str: String) : AbstractCoroutineContextElement(Key) {
        companion object Key : CoroutineContext.Key<CustomContextElement>
    }
}