package com.correx.core.tools

import com.correx.core.events.events.FileWrittenEvent
import com.correx.core.events.events.StoredEvent
import com.correx.core.events.events.ToolCallAssessedEvent
import com.correx.core.events.events.ToolExecutionCompletedEvent
import com.correx.core.events.events.ToolExecutionFailedEvent
import com.correx.core.events.events.ToolExecutionRejectedEvent
import com.correx.core.events.events.ToolExecutionStartedEvent
import com.correx.core.events.events.ToolInvocationRequestedEvent
import com.correx.core.events.types.ToolInvocationId
import com.correx.core.tools.state.FileMutationRecord
import com.correx.core.tools.state.ToolCallAssessmentRecord
import com.correx.core.tools.state.ToolInvocationRecord
import com.correx.core.tools.state.ToolInvocationStatus
import com.correx.core.tools.state.ToolState

class DefaultToolReducer : ToolReducer {
    override fun reduce(state: ToolState, event: StoredEvent): ToolState =
        when (val p = event.payload) {
            is ToolInvocationRequestedEvent -> state.copy(
                invocations = state.invocations + ToolInvocationRecord(
                    invocationId = p.invocationId,
                    toolName = p.toolName,
                    tier = p.tier,
                    request = p.request,
                    status = ToolInvocationStatus.REQUESTED,
                    requestedAt = event.metadata.timestamp
                )
            )
            is ToolExecutionStartedEvent -> state.updateRecord(p.invocationId) {
                it.copy(status = ToolInvocationStatus.STARTED)
            }
            is ToolExecutionCompletedEvent -> state.updateRecord(p.invocationId) {
                it.copy(
                    status = ToolInvocationStatus.COMPLETED,
                    receipt = p.receipt,
                    completedAt = event.metadata.timestamp
                )
            }
            is ToolExecutionFailedEvent -> state.updateRecord(p.invocationId) {
                it.copy(status = ToolInvocationStatus.FAILED, completedAt = event.metadata.timestamp)
            }
            is ToolExecutionRejectedEvent -> state.updateRecord(p.invocationId) {
                it.copy(status = ToolInvocationStatus.REJECTED, completedAt = event.metadata.timestamp)
            }
            is ToolCallAssessedEvent -> state.updateRecord(p.invocationId) {
                it.copy(
                    assessment = ToolCallAssessmentRecord(
                        issues = p.issues,
                        observations = p.observations,
                        disposition = p.disposition,
                    ),
                )
            }
            is FileWrittenEvent -> state.updateRecord(p.invocationId) {
                it.copy(
                    mutations = it.mutations + FileMutationRecord(
                        path = p.path,
                        preImageHash = p.preImageHash,
                        postImageHash = p.postImageHash,
                        preExisted = p.preExisted,
                        mode = p.mode,
                    ),
                )
            }
            else -> state
        }
}

private fun ToolState.updateRecord(
    invocationId: ToolInvocationId,
    transform: (ToolInvocationRecord) -> ToolInvocationRecord
): ToolState = copy(
    invocations = invocations.map { if (it.invocationId == invocationId) transform(it) else it }
)
