package com.correx.apps.server.bridge

import com.correx.apps.server.protocol.PauseReason
import com.correx.apps.server.protocol.ServerMessage
import com.correx.core.artifactstore.ArtifactStore
import com.correx.core.events.events.InferenceCompletedEvent
import com.correx.core.events.events.InferenceFailedEvent
import com.correx.core.events.events.InferenceStartedEvent
import com.correx.core.events.events.InferenceTimeoutEvent
import com.correx.core.events.events.OrchestrationPausedEvent
import com.correx.core.events.events.OrchestrationResumedEvent
import com.correx.core.events.events.RetryAttemptedEvent
import com.correx.core.events.events.StageCompletedEvent
import com.correx.core.events.events.StageFailedEvent
import com.correx.core.events.events.StoredEvent
import com.correx.core.events.events.TransitionExecutedEvent

/** Stage transition + inference lifecycle events (some carry timestamps / need the artifact store). */
@Suppress("LongMethod")
internal suspend fun mapStageInferenceEvent(
    event: StoredEvent,
    artifactStore: ArtifactStore,
    sessionSequence: Long,
): MapOutcome {
    val seq = event.sequence
    val msg = when (val p = event.payload) {
        is TransitionExecutedEvent -> ServerMessage.StageStarted(
            sessionId = p.sessionId,
            stageId = p.to,
            occurredAt = event.metadata.timestamp.toEpochMilliseconds(),
            sequence = seq,
            sessionSequence = sessionSequence,
        )

        is StageCompletedEvent -> ServerMessage.StageCompleted(
            sessionId = p.sessionId,
            stageId = p.stageId,
            occurredAt = event.metadata.timestamp.toEpochMilliseconds(),
            sequence = seq,
            sessionSequence = sessionSequence,
        )

        is StageFailedEvent -> ServerMessage.StageFailed(
            sessionId = p.sessionId,
            stageId = p.stageId,
            reason = p.reason,
            occurredAt = event.metadata.timestamp.toEpochMilliseconds(),
            sequence = seq,
            sessionSequence = sessionSequence,
        )

        is OrchestrationPausedEvent -> mapOrchestrationPaused(p, seq, sessionSequence)
        is OrchestrationResumedEvent -> ServerMessage.SessionResumed(
            sessionId = p.sessionId,
            stageId = p.stageId,
            sequence = seq,
            sessionSequence = sessionSequence,
        )

        is InferenceStartedEvent -> ServerMessage.InferenceStarted(
            sessionId = p.sessionId,
            stageId = p.stageId,
            sequence = seq,
            sessionSequence = sessionSequence,
        )

        is InferenceCompletedEvent -> mapInferenceCompleted(event, p, artifactStore, sessionSequence)
        is InferenceTimeoutEvent -> ServerMessage.InferenceTimedOut(
            sessionId = p.sessionId,
            stageId = p.stageId,
            elapsedMs = p.timeoutMs,
            sequence = seq,
            sessionSequence = sessionSequence,
        )

        is InferenceFailedEvent -> ServerMessage.InferenceFailed(
            sessionId = p.sessionId,
            stageId = p.stageId,
            reason = p.reason,
            sequence = seq,
            sessionSequence = sessionSequence,
        )

        is RetryAttemptedEvent -> ServerMessage.RetryAttempted(
            sessionId = p.sessionId,
            stageId = p.stageId,
            attemptNumber = p.attemptNumber,
            maxAttempts = p.maxAttempts,
            failureReason = p.failureReason,
            sequence = seq,
            sessionSequence = sessionSequence,
        )

        else -> return MapOutcome.Skip
    }
    return MapOutcome.Emit(msg)
}

private fun mapOrchestrationPaused(
    p: OrchestrationPausedEvent,
    seq: Long,
    sessionSequence: Long,
): ServerMessage {
    val reason = when (p.reason) {
        "APPROVAL_PENDING" -> PauseReason.APPROVAL_PENDING
        "CLARIFICATION_PENDING" -> PauseReason.CLARIFICATION_PENDING
        "ABANDONED_STALE" -> PauseReason.ABANDONED_STALE
        else -> PauseReason.USER_REQUESTED
    }
    return ServerMessage.SessionPaused(
        sessionId = p.sessionId,
        reason = reason,
        sequence = seq,
        sessionSequence = sessionSequence,
    )
}

private suspend fun mapInferenceCompleted(
    event: StoredEvent,
    p: InferenceCompletedEvent,
    artifactStore: ArtifactStore,
    sessionSequence: Long,
): ServerMessage {
    val response = runCatching {
        artifactStore.get(p.responseArtifactId)?.toString(Charsets.UTF_8) ?: ""
    }.getOrElse { "" }
    val reasoning = p.reasoningArtifactId?.let { id ->
        runCatching { artifactStore.get(id)?.toString(Charsets.UTF_8) ?: "" }.getOrElse { "" }
    } ?: ""
    return ServerMessage.InferenceCompleted(
        sessionId = p.sessionId,
        stageId = p.stageId,
        outputSummary = response,
        responseText = response,
        reasoning = reasoning,
        occurredAt = event.metadata.timestamp.toEpochMilliseconds(),
        totalTokens = p.tokensUsed.totalTokens,
        sequence = event.sequence,
        sessionSequence = sessionSequence,
    )
}
