import kotlinx.coroutines.CoroutineScope
import kotlinx.coroutines.Dispatchers
import kotlinx.coroutines.Job
import kotlinx.coroutines.SupervisorJob
import kotlinx.coroutines.async
import kotlinx.coroutines.awaitAll
import kotlinx.coroutines.cancel
import kotlinx.coroutines.coroutineScope
import kotlinx.coroutines.delay
import kotlinx.coroutines.launch
import kotlinx.coroutines.test.runTest
import org.junit.jupiter.api.Assertions.assertEquals
import org.junit.jupiter.api.Test
import java.util.concurrent.ConcurrentHashMap
import java.util.concurrent.atomic.AtomicInteger

/**
 * Validates that [ConcurrentHashMap.computeIfAbsent] combined with [CoroutineScope.launch]
 * prevents double execution of the launch body under concurrent access.
 *
 * This is the exact idiom used in [ServerModule.launchSessionRun] and
 * [ServerModule.launchSessionResume] to atomically register the running job
 * before launching it, replacing the non-atomic check-then-act pattern
 * (containsKey → launch → put) which allowed duplicate launches.
 */
class LaunchRegistrationRaceTest {

    @Test
    fun `concurrent computeIfAbsent only invokes mapping lambda once per key`(): Unit = runTest {
        val registry = ConcurrentHashMap<String, Job>()
        val launchCount = AtomicInteger(0)
        val scope = CoroutineScope(SupervisorJob() + Dispatchers.Default)
        val key = "session-1"

        coroutineScope {
            (1..50).map {
                async(Dispatchers.Default) {
                    registry.computeIfAbsent(key) { _: String ->
                        scope.launch {
                            launchCount.incrementAndGet()
                            delay(500)
                        }
                    }
                }
            }.awaitAll()
        }

        // computeIfAbsent returns the Job as soon as it's created; the launched body (which
        // increments launchCount) runs independently on `scope`. Join the registered jobs so the
        // increments have happened before we assert — otherwise launchCount races the launch bodies.
        registry.values.forEach { it.join() }
        assertEquals(1, launchCount.get(), "mapping lambda must execute exactly once")
        assertEquals(1, registry.size, "registry must contain exactly one entry")
        scope.cancel()
    }

    @Test
    fun `concurrent computeIfAbsent on different keys invokes each lambda independently`(): Unit = runTest {
        val registry = ConcurrentHashMap<String, Job>()
        val launchCount = AtomicInteger(0)
        val scope = CoroutineScope(SupervisorJob() + Dispatchers.Default)

        coroutineScope {
            (1..50).map { i ->
                async(Dispatchers.Default) {
                    registry.computeIfAbsent("session-$i") { _: String ->
                        scope.launch {
                            launchCount.incrementAndGet()
                            delay(100)
                        }
                    }
                }
            }.awaitAll()
        }

        // Join the launched bodies before asserting: computeIfAbsent returns the Job immediately,
        // but the increment runs on `scope` independently and may lag awaitAll (the 49-vs-50 flake).
        registry.values.forEach { it.join() }
        assertEquals(50, launchCount.get(), "each distinct key must execute its lambda")
        assertEquals(50, registry.size, "registry must contain 50 entries")
        scope.cancel()
    }
}
