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.protocol.TranscriptRowDto
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.InferenceCompletedEvent
import com.correx.core.events.events.InferenceStartedEvent
import com.correx.core.events.events.InferenceTimeoutEvent
import com.correx.core.events.events.ApprovalRequestedEvent
import com.correx.core.events.events.ArtifactCreatedEvent
import com.correx.core.events.events.ChatTurnEvent
import com.correx.core.events.events.ClarificationAnsweredEvent
import com.correx.core.events.events.ClarificationRequestedEvent
import com.correx.core.events.events.OrchestrationPausedEvent
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.TalkieNarrationEvent
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.stores.EventStore
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

// Pause reason emitted by SessionOrchestrator.requestClarificationIfNeeded. The reducer sets
// pendingApproval=true for every pause, so this reason string is the only reliable way to tell a
// clarification pause from an approval pause.
private const val CLARIFICATION_PENDING_REASON = "CLARIFICATION_PENDING"

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,
    // Request-ids of clarifications the orchestrator is still parked on with a live deferred. A
    // reconnecting client only re-opens a clarification modal that can actually be answered; after a
    // restart the deferred is gone until the parked stage is relaunched (see
    // ServerModule.resumeAbandonedSessions), which re-arms it.
    private val livePendingClarificationIds: () -> Set<String> = { emptySet() },
    private val send: suspend (ServerMessage) -> Unit,
) {
    private val approvalProjector = ApprovalProjector(DefaultApprovalReducer())
    private val mapper = DomainEventMapper(artifactStore)

    /**
     * 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()

            // NOTE: repairing stuck approval-pauses (PAUSED + pendingApproval but no unresolved
            // request, from a pre-Feb-13 orchestrator bug) used to happen HERE, per connection —
            // a write side effect on a pure read path that two concurrent clients could double-fire.
            // It now runs once at boot in ServerModule.repairStuckApprovalPauses(). replaySnapshot is
            // read-only again.

            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.
            pendingApprovalRequests.forEach { req ->
                approvalCoordinator?.registerPendingRequest(req.id, sessionId, req.tier)
            }

            val toolRecords = rebuildTools(events)
            val transcript = reconstructTranscript(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,
                transcript = transcript,
                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, tokenBudget = stageConfig.tokenBudget)
                }
                send(ServerMessage.StageToolManifest(
                    sessionId = sessionId,
                    workflowId = orchState.workflowId,
                    stages = stages,
                ))
            }

            rePushPendingClarification(orchState, events)
        }
        send(ServerMessage.SnapshotComplete)
    }

    /**
     * Re-surfaces a clarification the given session is parked on so a reconnecting client re-opens
     * the modal (parallel to the pendingApprovals re-registration above). Fires only for a
     * PAUSED/CLARIFICATION_PENDING session whose last unanswered ClarificationRequestedEvent still
     * has a *live* deferred in the orchestrator — a request the running process can actually resolve.
     * After a bare restart the deferred is gone until ServerModule.resumeAbandonedSessions relaunches
     * the parked stage (which re-emits a fresh, live ClarificationRequestedEvent); until then we do
     * not show a modal that would resume nothing.
     */
    private suspend fun rePushPendingClarification(
        orchState: com.correx.core.events.orchestration.OrchestrationState,
        events: List<StoredEvent>,
    ) {
        if (orchState.status != OrchestrationStatus.PAUSED) return
        if (orchState.pauseReason != CLARIFICATION_PENDING_REASON) return
        val answered = events.mapNotNullTo(mutableSetOf()) {
            (it.payload as? ClarificationAnsweredEvent)?.requestId?.value
        }
        val live = livePendingClarificationIds()
        events.lastOrNull {
            val p = it.payload
            p is ClarificationRequestedEvent && p.requestId.value !in answered && p.requestId.value in live
        }?.let { stored -> mapper.map(stored, stored.sessionSequence)?.let { send(it) } }
    }

    /**
     * Rebuild the OUTPUT-transcript source rows (chat / narration / thinking / tool) so a reopened
     * session restores its full conversation. Ships raw content only — the client applies its own
     * diff/param/action formatting. Reasoning is resolved from CAS (same as the live path). Tool
     * params are paired to their completion by invocation id.
     */
    private suspend fun reconstructTranscript(events: List<StoredEvent>): List<TranscriptRowDto> {
        val paramsByInvocation = mutableMapOf<String, List<String>>()
        val rows = mutableListOf<TranscriptRowDto>()
        for (event in events) {
            when (val p = event.payload) {
                is ChatTurnEvent -> rows.add(
                    TranscriptRowDto(
                        kind = "chat",
                        role = p.role.name,
                        content = p.content,
                        latencyMs = p.latencyMs,
                        totalTokens = p.tokensUsed?.totalTokens,
                    ),
                )
                is TalkieNarrationEvent -> rows.add(
                    TranscriptRowDto(
                        kind = "narration",
                        content = p.content,
                        latencyMs = p.latencyMs,
                        totalTokens = p.tokensUsed?.totalTokens,
                    ),
                )
                is InferenceCompletedEvent -> {
                    val reasoning = p.reasoningArtifactId
                        ?.let { runCatching { artifactStore.get(it)?.toString(Charsets.UTF_8) }.getOrNull() }
                        .orEmpty()
                    if (reasoning.isNotBlank()) {
                        rows.add(TranscriptRowDto(kind = "thinking", content = reasoning))
                    }
                }
                is ToolInvocationRequestedEvent ->
                    paramsByInvocation[p.invocationId.value] = prettyToolParams(p.request.parameters)
                is ToolExecutionCompletedEvent -> rows.add(
                    TranscriptRowDto(
                        kind = "tool",
                        status = "completed",
                        toolName = p.toolName,
                        diff = p.receipt.diff,
                        params = paramsByInvocation[p.invocationId.value].orEmpty(),
                        summary = p.receipt.outputSummary,
                        affectedEntities = p.receipt.affectedEntities,
                    ),
                )
                is ToolExecutionFailedEvent -> rows.add(
                    TranscriptRowDto(kind = "tool", status = "failed", toolName = p.toolName, summary = p.reason),
                )
                is ToolExecutionRejectedEvent -> rows.add(
                    TranscriptRowDto(kind = "tool", status = "rejected", toolName = p.toolName, summary = p.reason),
                )
                else -> { /* not a transcript row */ }
            }
        }
        return rows
    }

    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] }
    }

}
