import com.nomadgcs.data.event.EventPriority
import kotlinx.coroutines.CoroutineScope
import kotlinx.coroutines.Job
import kotlinx.coroutines.flow.MutableSharedFlow
import kotlinx.coroutines.flow.catch
import kotlinx.coroutines.flow.launchIn
import kotlinx.coroutines.flow.onEach
import kotlinx.coroutines.runBlocking
import kotlinx.coroutines.sync.Mutex
import kotlinx.coroutines.sync.withLock
import java.util.*
import kotlin.reflect.KClass

fun main(): Unit = runBlocking {
    // Create and register the handler
    val bus = EventRegistrar()
    var event: Any? = null
    val mutex = Mutex(true)
    bus.register(Any::class, EventPriority.NORMAL) {
        println("Called")
        event = it
        mutex.unlock()
    }
    // Call the event
    val testEvent = bus.emit(Any())
    println("Suspending")
    mutex.withLock {
        println("Acquired")
        println(testEvent == event)
    }
}

@Suppress("UNCHECKED_CAST")
class EventRegistrar(job: Job = Job()) : Job by job {

    private val scope = CoroutineScope(this)
    private val handlers = mutableMapOf<KClass<*>, EnumMap<EventPriority, MutableSharedFlow<*>>>()

    fun <T : Any> register(
        type: KClass<T>,
        priority: EventPriority,
        handler: suspend (T) -> Unit
    ): Job {
        println("Registering")
        val job = create(priority, type)
            .onEach(handler)
            .catch { println("Exception caught while handling event for ${type.simpleName}") }
            .launchIn(scope)
        println("Registered")
        return job
    }

    suspend fun <T : Any> emit(event: T): T {
        println("Emitting")
        val handlers = handlers[event::class] ?: return event
        // Since we're using an EnumMap, we know that values will always
        // be iterated in the order of the EventPriority declaration
        println("Emitting here")
//        println("Emitting to ${handlers.size}") // << Uncommenting will cause the problem to "fix" itself
        handlers.values.forEach {
            println("Subscribers")
            println("${it.subscriptionCount.value}")
//            println("Subscribers ${it.subscriptionCount.value}") // << Uncommenting will cause the problem to "fix" itself
            (it as MutableSharedFlow<T>).emit(event)
        }
        println("Done")
        return event
    }

    private fun <T : Any> create(priority: EventPriority, type: KClass<T>) =
        handlers.getOrPut(type) { EnumMap(EventPriority::class.java) }
            .getOrPut(priority) { MutableSharedFlow<T>() } as MutableSharedFlow<T>
}