package com.correx.apps.server.bridge

import com.correx.apps.server.approval.ApprovalCoordinator
import com.correx.apps.server.protocol.ApprovalDto
import com.correx.apps.server.protocol.EventEntryDto
import com.correx.apps.server.protocol.ServerMessage
import com.correx.apps.server.protocol.SessionStateDto
import com.correx.apps.server.protocol.StageToolDecl
import com.correx.apps.server.protocol.ToolDecl
import com.correx.apps.server.protocol.ToolRecordDto
import com.correx.apps.server.registry.WorkflowRegistry
import com.correx.core.approvals.ApprovalProjector
import com.correx.core.approvals.DefaultApprovalReducer
import com.correx.core.approvals.model.ApprovalState
import com.correx.core.artifactstore.ArtifactStore
import com.correx.core.events.events.EventMetadata
import com.correx.core.events.events.InferenceCompletedEvent
import com.correx.core.events.events.InferenceStartedEvent
import com.correx.core.events.events.InferenceTimeoutEvent
import com.correx.core.events.events.NewEvent
import com.correx.core.events.events.ApprovalRequestedEvent
import com.correx.core.events.events.ArtifactCreatedEvent
import com.correx.core.events.events.OrchestrationPausedEvent
import com.correx.core.events.events.OrchestrationResumedEvent
import com.correx.core.events.events.StageCompletedEvent
import com.correx.core.events.events.SessionWorkspaceBoundEvent
import com.correx.core.events.events.StageFailedEvent
import com.correx.core.events.events.StoredEvent
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.ToolInvocationRequestedEvent
import com.correx.core.events.events.TransitionExecutedEvent
import com.correx.core.events.events.WorkflowCompletedEvent
import com.correx.core.events.events.WorkflowStartedEvent
import com.correx.core.events.events.WorkflowFailedEvent
import com.correx.core.events.orchestration.OrchestrationStatus
import com.correx.core.events.types.EventId
import kotlinx.datetime.Clock
import com.correx.core.events.stores.EventStore
import com.correx.core.events.types.ApprovalRequestId
import com.correx.core.events.types.SessionId
import com.correx.core.kernel.orchestration.OrchestrationRepository
import com.correx.core.tools.registry.ToolRegistry

// How many of a session's events the reopen snapshot carries. The full log lives in the event
// store; this bounds the snapshot frame while being generous enough that a normal session's event
// list isn't truncated on reopen (was 7, which dropped most of a 100+ event run).
private const val SNAPSHOT_EVENT_LIMIT = 200

class SessionEventBridge(
    private val eventStore: EventStore,
    private val artifactStore: ArtifactStore,
    private val orchestrationRepository: OrchestrationRepository,
    private val workflowRegistry: WorkflowRegistry,
    private val toolRegistry: ToolRegistry,
    private val approvalCoordinator: ApprovalCoordinator? = null,
    private val send: suspend (ServerMessage) -> Unit,
) {
    private val approvalProjector = ApprovalProjector(DefaultApprovalReducer())

    /**
     * Resolves the last stage's inference output from CAS so a reopened session restores its output
     * panel. Returns empty if there's no completed inference or the artifact is unavailable.
     */
    private suspend fun resolveLastOutput(events: List<StoredEvent>): String =
        events.lastOrNull { it.payload is InferenceCompletedEvent }
            ?.let { (it.payload as InferenceCompletedEvent).responseArtifactId }
            ?.let { artifactStore.get(it)?.toString(Charsets.UTF_8) }
            .orEmpty()

    suspend fun replaySnapshot() {
        val lastGlobal = eventStore.lastGlobalSequence()
        eventStore.allSessionIds().forEach { sessionId ->
            val orchState = orchestrationRepository.getState(sessionId)
            if (orchState.status == OrchestrationStatus.IDLE) return@forEach

            val lastSession = eventStore.lastSequence(sessionId) ?: 0L

            // Single read from the event store — avoids inconsistency between
            // approval state and recentEvents with a cursor-based store.
            val events = eventStore.read(sessionId)

            val approvalState: ApprovalState? = if (orchState.pendingApproval) {
                events.fold(approvalProjector.initial()) { state: ApprovalState, event: StoredEvent ->
                    approvalProjector.apply(state, event)
                }
            } else null

            val pendingApprovalRequests = approvalState?.let { state ->
                state.requests.values
                    .filter { req -> state.decisions.values.none { it.requestId == req.id } }
                    .sortedWith(compareBy({ it.timestamp }, { it.id.value }))
            } ?: emptyList()

            // Fix stuck sessions: orchestration says PAUSED + pendingApproval but there are no
            // actual unresolved approval requests. This happens when approval was resolved but
            // OrchestrationResumedEvent was not emitted (pre-Feb-13 orchestrator bug). Append the
            // missing event permanently so the session state corrects on all future replays.
            //
            // Guard: if there are more OrchestrationPausedEvents than ApprovalRequestedEvents the
            // session just entered the approval gate and ApprovalRequestedEvent has not been stored
            // yet (tiny race window between two sequential emits). Do NOT fire in that case —
            // appending a spurious resume would hide the pending approval from the TUI.
            val pauseCount = events.count { it.payload is OrchestrationPausedEvent }
            val approvalRequestCount = events.count { it.payload is ApprovalRequestedEvent }
            val hasUnpairedPause = pauseCount > approvalRequestCount
            if (orchState.pendingApproval && pendingApprovalRequests.isEmpty() && !hasUnpairedPause) {
                val alreadyResumed = events.any { it.payload is OrchestrationResumedEvent }
                val stageId = orchState.currentStageId
                if (!alreadyResumed && stageId != null) {
                    val resumeEvent = NewEvent(
                        metadata = EventMetadata(
                            eventId = EventId(java.util.UUID.randomUUID().toString()),
                            sessionId = sessionId,
                            timestamp = Clock.System.now(),
                            schemaVersion = 1,
                            causationId = null,
                            correlationId = null,
                        ),
                        payload = OrchestrationResumedEvent(
                            sessionId = sessionId,
                            stageId = stageId,
                        ),
                    )
                    eventStore.append(resumeEvent)
                }
            }

            val pendingApprovals = pendingApprovalRequests.map {
                ApprovalDto(
                    requestId = it.id.value,
                    tier = it.tier.name,
                    toolName = it.toolName,
                    preview = it.preview,
                )
            }

            val recentEvents = events
                .mapNotNull { eventToEntry(it) }
                .takeLast(SNAPSHOT_EVENT_LIMIT)

            // Restore the output panel on reopen (otherwise the snapshot carries no output).
            val lastOutput = resolveLastOutput(events)
            val workspaceRoot = events.lastOrNull { it.payload is SessionWorkspaceBoundEvent }
                ?.let { (it.payload as SessionWorkspaceBoundEvent).workspaceRoot }
                .orEmpty()

            // Re-register pending approvals so the ApprovalCoordinator can route responses
            // from clients that connected after the ApprovalRequestedEvent was emitted.
            pendingApprovals.forEach { dto ->
                approvalCoordinator?.registerPendingRequest(ApprovalRequestId(dto.requestId), sessionId)
            }

            val toolRecords = rebuildTools(events)

            send(ServerMessage.SessionSnapshot(
                sessionId = sessionId,
                workflowId = orchState.workflowId,
                state = SessionStateDto(
                    status = orchState.status.name,
                    currentStageId = orchState.currentStageId?.value,
                    pauseReason = orchState.pauseReason,
                ),
                pendingApprovals = pendingApprovals,
                tools = toolRecords,
                recentEvents = recentEvents,
                lastOutput = lastOutput,
                lastResponse = lastOutput,
                workspaceRoot = workspaceRoot,
                lastSequence = lastGlobal,
                lastSessionSequence = lastSession,
            ))

            // Emit StageToolManifest for this session (ARCH-TOOL-1)
            val graph = workflowRegistry.find(orchState.workflowId)
            if (graph != null) {
                val stages = graph.stages.map { (stageId, stageConfig) ->
                    val tools = stageConfig.allowedTools.mapNotNull { toolName ->
                        toolRegistry.resolve(toolName)?.let { tool ->
                            ToolDecl(name = tool.name, tier = tool.tier.level)
                        }
                    }
                    StageToolDecl(stageId = stageId.value, tools = tools)
                }
                send(ServerMessage.StageToolManifest(
                    sessionId = sessionId,
                    workflowId = orchState.workflowId,
                    stages = stages,
                ))
            }
        }
        send(ServerMessage.SnapshotComplete)
    }

    private fun eventToEntry(event: StoredEvent): EventEntryDto? {
        val ts = event.metadata.timestamp.toEpochMilliseconds()
        return when (val p = event.payload) {
            is TransitionExecutedEvent -> EventEntryDto(ts, "StageStarted", p.to.value)
            is StageCompletedEvent -> EventEntryDto(ts, "StageCompleted", p.stageId.value)
            is StageFailedEvent -> EventEntryDto(ts, "StageFailed", p.stageId.value)
            is InferenceStartedEvent -> EventEntryDto(ts, "InferenceStarted", p.stageId.value)
            is InferenceCompletedEvent -> EventEntryDto(ts, "InferenceCompleted", p.stageId.value)
            is InferenceTimeoutEvent -> EventEntryDto(ts, "InferenceTimeout", p.stageId.value)
            is ToolInvocationRequestedEvent -> EventEntryDto(ts, "ToolStarted", p.toolName)
            is ToolExecutionCompletedEvent -> EventEntryDto(ts, "ToolCompleted", p.toolName)
            is ToolExecutionFailedEvent -> EventEntryDto(ts, "ToolFailed", p.toolName)
            is ToolExecutionRejectedEvent -> EventEntryDto(ts, "ToolRejected", p.toolName)
            is OrchestrationPausedEvent -> EventEntryDto(ts, "SessionPaused", p.reason)
            is WorkflowCompletedEvent -> EventEntryDto(ts, "SessionCompleted", "")
            is WorkflowFailedEvent -> EventEntryDto(ts, "SessionFailed", p.reason)
            is ArtifactCreatedEvent -> EventEntryDto(ts, "ArtifactCreated", p.artifactId.value)
            else -> null
        }
    }

    /**
     * Rebuild tool execution records from stored events in order.
     * Tracks tools by invocation ID across stage transitions — each stage transition clears
     * the tool list so the snapshot only carries tools from the current stage.
     * Result preserves insertion order (oldest invocation first).
     */
    private fun rebuildTools(events: List<StoredEvent>): List<ToolRecordDto> {
        val orderedIds = mutableListOf<String>()
        val byInvocation = mutableMapOf<String, ToolRecordDto>()

        for (event in events) {
            when (val p = event.payload) {
                is WorkflowStartedEvent -> {
                    orderedIds.clear(); byInvocation.clear()
                }
                is TransitionExecutedEvent -> {
                    orderedIds.clear(); byInvocation.clear()
                }
                is ToolInvocationRequestedEvent -> {
                    val id = p.invocationId.value
                    if (!byInvocation.containsKey(id)) {
                        orderedIds.add(id)
                    }
                    byInvocation[id] = ToolRecordDto(
                        name = p.toolName,
                        tier = p.tier.level,
                        status = "STARTED",
                    )
                }
                is ToolExecutionCompletedEvent -> {
                    val id = p.invocationId.value
                    byInvocation[id]?.let { existing ->
                        byInvocation[id] = existing.copy(
                            status = "COMPLETED",
                            diff = p.receipt.diff,
                        )
                    }
                }
                is ToolExecutionFailedEvent -> {
                    val id = p.invocationId.value
                    byInvocation[id]?.let { existing ->
                        byInvocation[id] = existing.copy(status = "FAILED")
                    }
                }
                is ToolExecutionRejectedEvent -> {
                    val id = p.invocationId.value
                    byInvocation[id]?.let { existing ->
                        byInvocation[id] = existing.copy(status = "REJECTED")
                    }
                }
                else -> { /* skip — irrelevant event types */ }
            }
        }

        return orderedIds.mapNotNull { byInvocation[it] }
    }

}
