package com.correx.core.inference

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.StoredEvent
import com.correx.core.events.types.InferenceRequestId

class DefaultInferenceReducer : InferenceReducer {
    override fun reduce(state: InferenceState, event: StoredEvent): InferenceState = when (val p = event.payload) {
        is InferenceStartedEvent -> {
            val newRecord = InferenceRecord(
                requestId = p.requestId,
                sessionId = p.sessionId,
                stageId = p.stageId,
                providerId = p.providerId,
                status = InferenceStatus.STARTED,
            )
            state.copy(records = state.records + newRecord)
        }

        is InferenceCompletedEvent -> updateRecord(state, p.requestId) {
            it.copy(status = InferenceStatus.COMPLETED, tokensUsed = p.tokensUsed, latencyMs = p.latencyMs)
        }

        is InferenceFailedEvent -> updateRecord(state, p.requestId) {
            it.copy(status = InferenceStatus.FAILED, failureReason = p.reason)
        }

        is InferenceTimeoutEvent -> updateRecord(state, p.requestId) {
            it.copy(status = InferenceStatus.TIMED_OUT, latencyMs = p.timeoutMs)
        }

        else -> state
    }

    private fun updateRecord(
        state: InferenceState,
        requestId: InferenceRequestId,
        transform: (InferenceRecord) -> InferenceRecord,
    ): InferenceState {
        val updated = state.records.map { if (it.requestId == requestId) transform(it) else it }
        return state.copy(records = updated)
    }

}
