package com.correx.apps.server.logging

import com.correx.core.events.types.SessionId
import com.correx.core.events.types.StageId
import kotlinx.coroutines.slf4j.MDCContext
import kotlinx.coroutines.withContext
import org.slf4j.MDC

/**
 * Execute [block] with [sessionId] and optional [stageId] bound to the SLF4J MDC (Mapped Diagnostic Context).
 *
 * The sessionId and stageId are automatically propagated across coroutine suspension points via
 * [MDCContext] (requires kotlinx-coroutines-slf4j). Prior MDC values are saved and restored after
 * the block exits, enabling safe nesting.
 *
 * @param sessionId The session identifier to bind in MDC under key "sessionId"
 * @param stageId Optional stage identifier to bind in MDC under key "stageId"; if null, any existing
 *                 stageId is cleared
 * @param block The suspending function to execute with MDC context active
 * @return The result of [block]
 */
suspend inline fun <T> withSessionContext(
    sessionId: SessionId,
    stageId: StageId? = null,
    crossinline block: suspend () -> T,
): T {
    // Save prior values so we can restore them after the block.
    val priorSessionId = MDC.get("sessionId")
    val priorStageId = MDC.get("stageId")

    return try {
        // Set the new values in MDC.
        MDC.put("sessionId", sessionId.value)
        if (stageId != null) {
            MDC.put("stageId", stageId.value)
        } else {
            // Clear stageId if not provided, in case one was set by an outer context.
            MDC.remove("stageId")
        }

        // Wrap the block in MDCContext to propagate MDC across coroutine boundaries.
        // MDCContext captures the current MDC state and restores it when the coroutine resumes
        // after a suspension point.
        withContext(MDCContext()) {
            block()
        }
    } finally {
        // Restore prior values. If priorSessionId was null, remove the key entirely.
        if (priorSessionId != null) {
            MDC.put("sessionId", priorSessionId)
        } else {
            MDC.remove("sessionId")
        }
        if (priorStageId != null) {
            MDC.put("stageId", priorStageId)
        } else {
            MDC.remove("stageId")
        }
    }
}
