package com.correx.apps.server.ws

import com.correx.apps.server.ServerModule
import com.correx.apps.server.bridge.DomainEventMapper
import com.correx.apps.server.logging.withSessionContext
import com.correx.apps.server.protocol.ClientMessage
import com.correx.apps.server.protocol.ProtocolSerializer
import com.correx.apps.server.protocol.ServerMessage
import com.correx.apps.server.protocol.ServerMessage.TalkieResponseMessage
import com.correx.core.events.types.SessionId
import io.ktor.server.websocket.DefaultWebSocketServerSession
import io.ktor.websocket.Frame
import io.ktor.websocket.readText
import kotlinx.coroutines.channels.ClosedReceiveChannelException
import org.slf4j.LoggerFactory

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

class SessionStreamHandler(private val module: ServerModule) {

    private val mapper = DomainEventMapper(module.artifactStore)

    suspend fun handle(session: DefaultWebSocketServerSession, sessionId: SessionId, lastEventId: Long?) {
        log.info("client connected session={} lastEventId={}", sessionId.value, lastEventId)

        val replayEvents = if (lastEventId != null) {
            module.eventStore.readFrom(sessionId, lastEventId)
        } else {
            emptyList()
        }
        log.debug("replaying {} event(s) for session={}", replayEvents.size, sessionId.value)

        for (event in replayEvents) {
            runCatching { mapper.map(event) }
                .onFailure { log.warn("skipping unmappable event seq={}: {}", event.sequence, it.message) }
                .getOrNull()
                ?.let { msg ->
                    session.send(Frame.Text(ProtocolSerializer.encodeServerMessage(msg)))
                }
        }

        module.approvalCoordinator.registerClient(sessionId, session)
        try {
            for (frame in session.incoming) {
                if (frame is Frame.Text) {
                    val text = frame.readText()
                    runCatching { ProtocolSerializer.decodeClientMessage(text) }
                        .onSuccess { msg ->
                            log.debug("recv: {}", msg::class.simpleName)
                            handleClientMessage(session, sessionId, msg)
                        }
                        .onFailure {
                            log.warn("decode error: {}", it.message)
                            val error = ServerMessage.ProtocolError("Unknown message: ${it.message}")
                            session.send(Frame.Text(ProtocolSerializer.encodeServerMessage(error)))
                        }
                }
            }
        } catch (_: ClosedReceiveChannelException) {
            log.info("client disconnected session={}", sessionId.value)
        } finally {
            module.approvalCoordinator.unregisterClient(sessionId, session)
        }
    }

    private suspend fun handleClientMessage(
        session: DefaultWebSocketServerSession,
        sessionId: SessionId,
        msg: ClientMessage,
    ) {
        when (msg) {
            is ClientMessage.Ping -> Unit
            is ClientMessage.CancelSession -> {
                runCatching { module.orchestrator.cancel(msg.sessionId) }
                    .onFailure { log.error("cancel failed for session={}: {}", msg.sessionId.value, it.message, it) }
            }

            is ClientMessage.ApprovalResponse -> {
                module.approvalCoordinator.handleResponse(msg, sessionId)?.let { reply ->
                    session.send(Frame.Text(ProtocolSerializer.encodeServerMessage(reply)))
                }
            }

            is ClientMessage.ChatInput -> {
                withSessionContext(msg.sessionId) {
                    runCatching {
                        module.routerFacade.onUserInput(
                            sessionId = msg.sessionId,
                            input = msg.text,
                            mode = msg.mode,
                        )
                    }.onSuccess { response ->
                        session.send(
                            Frame.Text(
                                ProtocolSerializer.encodeServerMessage(
                                    TalkieResponseMessage(
                                        sessionId = msg.sessionId,
                                        content = response.content,
                                        steeringEmitted = response.steeringEmitted,
                                    ),
                                ),
                            ),
                        )
                    }.onFailure {
                        log.error(
                            "routerFacade.onUserInput failed for session={}: {}",
                            msg.sessionId.value,
                            it.message,
                            it,
                        )
                        val error = ServerMessage.ProtocolError("Router error: ${it.message}")
                        session.send(Frame.Text(ProtocolSerializer.encodeServerMessage(error)))
                    }
                }
            }

            is ClientMessage.ClarificationResponse -> {
                runCatching {
                    module.orchestrator.submitClarification(msg.sessionId, msg.stageId, msg.requestId, msg.answers)
                }.onFailure {
                    log.error("submitClarification failed for session={}: {}", msg.sessionId.value, it.message, it)
                }
            }

            else -> {
                log.warn("unexpected message type={} in session stream", msg::class.simpleName)
                val error = ServerMessage.ProtocolError("Unexpected message type in session stream")
                session.send(Frame.Text(ProtocolSerializer.encodeServerMessage(error)))
            }
        }
    }

}
