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.RetrySalvageDecidedEvent
import com.correx.core.events.events.SalvageDecision
import com.correx.core.events.events.FailureTicketOpenedEvent
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

// Budget-key suffix distinguishing the tier-2 escalation (arbiter) budget from the tier-1 owner
// budget for the same failing stage. Shared with DefaultSessionOrchestrator (same package).
internal const val INTENT_BUDGET_SUFFIX = "#intent"

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,
        )

        // Entering a new stage resets the retry budget: retryCount is per-stage (RetryAttemptedEvent
        // is stage-scoped, maxAttempts is a per-stage cap), so without this reset the budget is shared
        // session-wide — an early stage that exhausts its retries starves every later stage of its own,
        // turning the next stage's first retryable failure into a terminal WorkflowFailed (live repro:
        // implement burned 3/3, then review's retryable JSON-decode error got 0 retries). Back-edge
        // loops stay bounded by the separate refinementIterations counter, so this can't loop forever.
        is TransitionExecutedEvent -> state.copy(
            currentStageId = p.to,
            retryCount = 0,
            gateRetryBudgets = emptyMap(),
            gateFailureFingerprints = emptyMap(),
            gateSalvageUsed = emptySet(),
        )

        is RetryAttemptedEvent -> state.copy(
            status = OrchestrationStatus.RUNNING,
            retryCount = p.attemptNumber,
            failureReason = null,
            gateRetryBudgets = if (p.charged) {
                state.gateRetryBudgets + (p.gate to ((state.gateRetryBudgets[p.gate] ?: 0) + 1))
            } else {
                state.gateRetryBudgets
            },
            gateFailureFingerprints = state.gateFailureFingerprints + (p.gate to p.fingerprint),
        )

        is RetrySalvageDecidedEvent -> if (p.decision == SalvageDecision.CONTINUE) {
            state.copy(
                gateRetryBudgets = state.gateRetryBudgets + (p.gate to 0),
                gateSalvageUsed = state.gateSalvageUsed + p.gate,
            )
        } else {
            state
        }

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

        // Charge the failing stage's recovery-route budget. routeAttempt is the no-progress count the
        // orchestrator computed (progress-aware), so fold it directly (idempotent under replay), and
        // record the failure fingerprint the next route compares against. Keyed per TIER: the tier-1
        // owner loop and the tier-2 escalation (arbiter) hold independent budgets for the same failing
        // stage, so an escalated ticket ("#intent" suffix) never spends the owner budget and vice versa.
        is FailureTicketOpenedEvent -> {
            val budgetKey = p.stageId.value + if (p.escalated) INTENT_BUDGET_SUFFIX else ""
            state.copy(
                recoveryRoutes = state.recoveryRoutes + (budgetKey to p.routeAttempt),
                recoveryFailureFingerprints = state.recoveryFailureFingerprints + (budgetKey to p.fingerprint),
            )
        }

        else -> state
    }
}
