package com.correx.apps.server.approval

import com.correx.apps.server.protocol.ClientMessage
import com.correx.apps.server.protocol.ProtocolSerializer
import com.correx.apps.server.protocol.RiskSummaryDto
import com.correx.apps.server.protocol.ServerMessage
import com.correx.apps.server.protocol.toDto
import com.correx.core.approvals.Tier
import com.correx.core.events.events.ApprovalRequestedEvent
import com.correx.core.events.risk.RiskAction
import com.correx.core.events.types.ApprovalRequestId
import com.correx.core.events.types.SessionId
import com.correx.core.kernel.orchestration.ApprovalGateway
import io.ktor.server.websocket.DefaultWebSocketServerSession
import io.ktor.websocket.Frame
import org.slf4j.LoggerFactory
import java.util.concurrent.ConcurrentHashMap

open class ApprovalCoordinator(
    private val orchestrator: ApprovalGateway,
) {
    private val sessionClients: ConcurrentHashMap<SessionId, MutableSet<DefaultWebSocketServerSession>> =
        ConcurrentHashMap()
    private val globalClients: MutableSet<DefaultWebSocketServerSession> = ConcurrentHashMap.newKeySet()
    private val resolved: ConcurrentHashMap<ApprovalRequestId, Boolean> = ConcurrentHashMap()
    private val requestSessions: ConcurrentHashMap<ApprovalRequestId, SessionId> = ConcurrentHashMap()
    private val requestTiers: ConcurrentHashMap<ApprovalRequestId, Tier> = ConcurrentHashMap()

    fun registerClient(sessionId: SessionId, session: DefaultWebSocketServerSession) {
        sessionClients.getOrPut(sessionId) { ConcurrentHashMap.newKeySet() }.add(session)
    }

    fun unregisterClient(sessionId: SessionId, session: DefaultWebSocketServerSession) {
        sessionClients[sessionId]?.remove(session)
    }

    fun registerGlobalClient(session: DefaultWebSocketServerSession) {
        globalClients.add(session)
    }

    fun unregisterGlobalClient(session: DefaultWebSocketServerSession) {
        globalClients.remove(session)
    }

    suspend fun onApprovalRequested(event: ApprovalRequestedEvent, sessionSequence: Long = 0L) {
        requestSessions[event.requestId] = event.sessionId
        requestTiers[event.requestId] = event.tier
        val msg = ServerMessage.ApprovalRequired(
            sessionId = event.sessionId,
            requestId = event.requestId,
            tier = event.tier,
            riskSummary = event.riskSummary?.toDto() ?: RiskSummaryDto(
                level = event.tier.name,
                factors = emptyList(),
                recommendedAction = RiskAction.PROMPT_USER.name,
                rationale = emptyList(),
            ),
            toolName = event.toolName,
            preview = event.preview,
            sequence = 0L,
            sessionSequence = sessionSequence,
        )
        broadcast(event.sessionId, msg)
    }

    fun lookupSession(requestId: ApprovalRequestId): SessionId? = requestSessions[requestId]

    /**
     * The pending approval request for [sessionId], if any. A session suspends on at most one
     * approval at a time, so this reverse lookup is unambiguous. Used by the REST approve route,
     * which is keyed by session (the caller need not know the requestId).
     */
    fun pendingRequestFor(sessionId: SessionId): ApprovalRequestId? =
        requestSessions.entries.firstOrNull { it.value == sessionId }?.key

    /**
     * Register a pending approval request that was reconstructed from the persisted event store
     * during snapshot replay. This makes [lookupSession] work for approvals that were created
     * before the current connection's live stream started.
     */
    fun registerPendingRequest(requestId: ApprovalRequestId, sessionId: SessionId, tier: Tier) {
        requestSessions[requestId] = sessionId
        requestTiers[requestId] = tier
    }

    open suspend fun handleResponse(msg: ClientMessage.ApprovalResponse, sessionId: SessionId): ServerMessage? {
        if (resolved.putIfAbsent(msg.requestId, true) != null) {
            return ServerMessage.ProtocolError(
                message = "Approval request ${msg.requestId.value} was already resolved",
                sequence = null,
                sessionSequence = null,
            )
        }
        val tier = requestTiers[msg.requestId] ?: Tier.T2
        val domain = msg.toDomain(sessionId, null, tier)
        return runCatching { orchestrator.submitApprovalDecision(msg.requestId, domain) }
            .fold(
                onSuccess = {
                    // Only retire the request once the decision was actually recorded. If submit
                    // throws, we clear the resolved flag so the client can retry instead of the
                    // request being permanently unanswerable.
                    requestSessions.remove(msg.requestId)
                    requestTiers.remove(msg.requestId)
                    null
                },
                onFailure = {
                    resolved.remove(msg.requestId)
                    ServerMessage.ProtocolError(
                        message = it.message ?: "Unknown error",
                        sequence = null,
                        sessionSequence = null,
                    )
                },
            )
    }

    private suspend fun broadcast(sessionId: SessionId, msg: ServerMessage) {
        val encoded = ProtocolSerializer.encodeServerMessage(msg)
        val sessionSubs = sessionClients[sessionId].orEmpty()
        val recipients: Set<DefaultWebSocketServerSession> = sessionSubs + globalClients
        var delivered = 0
        recipients.forEach { client ->
            runCatching { client.send(Frame.Text(encoded)) }
                .onSuccess { delivered++ }
                .onFailure { err ->
                    // A dead socket must not silently swallow an approval prompt. Log it and evict
                    // the client from both registries so we don't keep sending into a closed socket
                    // (the client re-registers and gets the pending approval from the snapshot on
                    // reconnect).
                    log.warn("approval broadcast to a client failed, evicting: {}", err.message)
                    sessionClients[sessionId]?.remove(client)
                    globalClients.remove(client)
                }
        }
        if (delivered == 0) {
            log.error(
                "approval prompt for session={} reached 0 live clients; session stays paused until a client reconnects",
                sessionId.value,
            )
        }
    }

    private companion object {
        private val log = LoggerFactory.getLogger(ApprovalCoordinator::class.java)
    }
}
