Async-Spring-Kotlin-JOOQ-Demo/src/main/kotlin/com/example/demo/controller/RsocketConnectionController.kt

package com.example.demo.controller

import com.example.demo.model.Message
import com.example.demo.model.User
import io.vavr.collection.HashMap
import io.vavr.collection.Map
import kotlinx.coroutines.delay
import kotlinx.coroutines.flow.Flow
import kotlinx.coroutines.flow.flow
import org.slf4j.LoggerFactory
import org.springframework.messaging.handler.annotation.MessageExceptionHandler
import org.springframework.messaging.handler.annotation.MessageMapping
import org.springframework.messaging.handler.annotation.Payload
import org.springframework.messaging.rsocket.RSocketRequester
import org.springframework.messaging.rsocket.annotation.ConnectMapping
import org.springframework.stereotype.Controller

@Controller
class RSocketConnectionController {

    private val log = LoggerFactory.getLogger(RSocketConnectionController::class.java)

    private var requesterMap: Map<String, RSocketRequester> = HashMap.empty()

    @Synchronized
    private fun getRequesterMap(): Map<String, RSocketRequester> {
        return requesterMap
    }

    @Synchronized
    private fun addRequester(rSocketRequester: RSocketRequester, clientId: String) {
        log.info("adding requester {}", clientId)
        requesterMap = requesterMap.put(clientId, rSocketRequester)
    }

    @Synchronized
    private fun removeRequester(clientId: String) {
        log.info("removing requester {}", clientId)
        requesterMap = requesterMap.remove(clientId)
    }

    @ConnectMapping("client-id")
    fun onConnect(rSocketRequester: RSocketRequester, @Payload user: User) {
        //        val clientIdFixed = clientId.replace("\"", "")  //check why the serializer  adds " to strings
//        rSocketRequester.rsocket().dispose()   //to reject connection
        val clientId = user.id.toString()
        rSocketRequester
                .rsocket()
                .onClose()
                .subscribe(null, null, {
                    log.info("{} just disconnected", clientId)
                    removeRequester(clientId)
                })
        addRequester(rSocketRequester, clientId)
    }

    @MessageMapping("private.news")
    fun privateNews(message: Message, rSocketRequesterParam: RSocketRequester) {
        getRequesterMap()
                .filterKeys { key -> key == message.toUser || key == message.fromUser }
                .values()
                .forEach { requester ->
                    run {
                        println("Sending message")
                        sendMessage(requester, message)
                    }
                }

    }

    @MessageExceptionHandler
    suspend fun handleException(ex: IllegalArgumentException): String {
        delay(10)
        return "${ex.message} handled"
    }

    @MessageMapping("echo-stream-async")
    suspend fun echoStreamAsync(payload: String): Flow<String> {
        delay(10)
        var i = 0
        return flow {
            while (true) {
                delay(10)
                emit("$payload ${i++}")
            }
        }
    }

    private fun sendMessage(requester: RSocketRequester, message: Message) =
            requester
                    .route("user.queue.reply")
//                    .metadata("test", MimeTypeUtils.parseMimeType(WellKnownMimeType.MESSAGE_RSOCKET_COMPOSITE_METADATA.string))
                    .data(message)
                    .send()
                    .subscribe()

}