package com.correx.core.kernel.orchestration

import com.correx.core.events.events.ChatSessionStartedEvent
import com.correx.core.events.events.OrchestrationPausedEvent
import com.correx.core.events.events.OrchestrationResumedEvent
import com.correx.core.events.events.RefinementIterationEvent
import com.correx.core.events.events.RetryAttemptedEvent
import com.correx.core.events.events.StoredEvent
import com.correx.core.events.events.TransitionExecutedEvent
import com.correx.core.events.events.WorkflowCompletedEvent
import com.correx.core.events.events.WorkflowFailedEvent
import com.correx.core.events.events.WorkflowStartedEvent
import com.correx.core.events.orchestration.OrchestrationState
import com.correx.core.events.orchestration.OrchestrationStatus

class DefaultOrchestrationReducer : OrchestrationReducer {
    override fun reduce(
        state: OrchestrationState,
        event: StoredEvent,
    ): OrchestrationState = when (val p = event.payload) {
        is ChatSessionStartedEvent -> state.copy(
            workflowId = "chat",
            status = OrchestrationStatus.CHAT,
            currentStageId = null,
        )

        is WorkflowStartedEvent -> state.copy(
            workflowId = p.workflowId,
            status = OrchestrationStatus.RUNNING,
            currentStageId = p.startStageId,
            retryPolicy = p.retryPolicy,
        )

        is WorkflowCompletedEvent -> state.copy(
            status = OrchestrationStatus.COMPLETED,
            currentStageId = p.terminalStageId,
        )

        is WorkflowFailedEvent -> state.copy(
            status = OrchestrationStatus.FAILED,
            failureReason = p.reason,
        )

        is OrchestrationPausedEvent -> state.copy(
            status = OrchestrationStatus.PAUSED,
            pendingApproval = true,
            pauseReason = p.reason,
        )

        is OrchestrationResumedEvent -> state.copy(
            status = OrchestrationStatus.RUNNING,
            pendingApproval = false,
            pauseReason = null,
        )

        is TransitionExecutedEvent -> state.copy(currentStageId = p.to)

        is RetryAttemptedEvent -> state.copy(
            status = OrchestrationStatus.RUNNING,
            retryCount = p.attemptNumber,
            failureReason = null,
        )

        is RefinementIterationEvent -> state.copy(
            refinementIterations = state.refinementIterations + (p.cycleKey to p.iteration),
        )

        else -> state
    }
}
