package com.correx.core.router

import com.correx.core.events.events.ChatTurnEvent
import com.correx.core.events.events.ChatTurnRole
import com.correx.core.events.events.ContextTruncatedEvent
import com.correx.core.events.events.EventMetadata
import com.correx.core.events.events.L3MemoryRetrievedEvent
import com.correx.core.events.events.L3RetrievedHit
import com.correx.core.events.events.NewEvent
import com.correx.core.events.events.RouterNarrationEvent
import com.correx.core.events.events.SteeringNoteAddedEvent
import com.correx.core.events.stores.EventStore
import com.correx.core.events.types.EventId
import com.correx.core.events.types.InferenceRequestId
import com.correx.core.events.types.SessionId
import com.correx.core.events.types.StageId
import com.correx.core.inference.Embedder
import com.correx.core.inference.InferenceRequest
import com.correx.core.inference.InferenceRouter
import com.correx.core.inference.ModelCapability
import com.correx.core.inference.ResponseFormat
import com.correx.core.inference.TokenUsage
import com.correx.core.router.l3.L3MemoryEntry
import com.correx.core.router.l3.L3MemoryStore
import com.correx.core.router.l3.L3Query
import com.correx.core.router.model.NarrationTrigger
import com.correx.core.router.model.RouterConfig
import com.correx.core.router.model.RouterResponse
import kotlinx.coroutines.CancellationException
import kotlinx.datetime.Clock
import org.slf4j.LoggerFactory
import java.util.UUID

private val log = LoggerFactory.getLogger(DefaultRouterFacade::class.java)

interface RouterFacade {
    suspend fun onUserInput(
        sessionId: SessionId,
        input: String,
        mode: ChatMode = ChatMode.CHAT,
    ): RouterResponse

    suspend fun narrate(sessionId: SessionId, trigger: NarrationTrigger)
}

class DefaultRouterFacade(
    private val routerRepository: RouterRepository,
    private val routerContextBuilder: RouterContextBuilder,
    private val inferenceRouter: InferenceRouter,
    private val eventStore: EventStore,
    private val config: RouterConfig,
    private val validateSteering: (suspend (String) -> String?)? = null,
    private val embedder: Embedder,
    private val l3MemoryStore: L3MemoryStore,
) : RouterFacade {

    override suspend fun onUserInput(
        sessionId: SessionId,
        input: String,
        mode: ChatMode,
    ): RouterResponse {
        // Emit USER turn event and get the turnId for retrieval dedup
        val userTurnId = emitChatTurn(sessionId, input, ChatTurnRole.USER)

        // Rebuild state with user turn appended
        var state = routerRepository.getRouterState(sessionId)
        val effectiveStageId = state.currentStageId ?: StageId.NONE

        // L3 retrieval — non-fatal; results fed back into state via event.
        // Always emit L3MemoryRetrievedEvent for every CHAT turn where retrieval was attempted,
        // including with empty hits — so lastRetrievedMemory is always current-turn-scoped.
        val retrieved = runCatching {
            val queryVector = embedder.embed(input)
            l3MemoryStore.query(L3Query(vector = queryVector, k = config.retrievalK))
        }.also { result ->
            result.exceptionOrNull()?.let { e ->
                if (e is CancellationException) throw e
                log.warn("L3 retrieval failed for session {}: {}", sessionId, e.message)
            }
        }.getOrElse { emptyList() }

        val inSessionTurnIds = state.conversationHistory.map { it.turnId }.toSet()
        val deduped = retrieved.filter { it.entry.turnId !in inSessionTurnIds }

        // Always emit L3MemoryRetrievedEvent (even with empty hits) so that
        // lastRetrievedMemory always reflects this turn and a prior turn's hits
        // are never injected into the next turn's context.
        val retrievalNow = Clock.System.now()
        eventStore.append(
            NewEvent(
                metadata = EventMetadata(
                    eventId = EventId(UUID.randomUUID().toString()),
                    sessionId = sessionId,
                    timestamp = retrievalNow,
                    schemaVersion = 1,
                    causationId = null,
                    correlationId = null,
                ),
                payload = L3MemoryRetrievedEvent(
                    sessionId = sessionId,
                    queryTurnId = userTurnId,
                    hits = deduped.map { hit ->
                        L3RetrievedHit(
                            entryId = hit.entry.id,
                            sourceSessionId = hit.entry.sessionId,
                            sourceTurnId = hit.entry.turnId,
                            text = hit.entry.text,
                            score = hit.score,
                        )
                    },
                    timestampMs = retrievalNow.toEpochMilliseconds(),
                ),
            ),
        )
        // Rebuild state so lastRetrievedMemory reflects this turn
        state = routerRepository.getRouterState(sessionId)

        val contextPack = routerContextBuilder.build(state, config.tokenBudget)

        if (contextPack.compressionMetadata.entriesDropped > 0) {
            val truncNow = Clock.System.now()
            eventStore.append(
                NewEvent(
                    metadata = EventMetadata(
                        eventId = EventId(UUID.randomUUID().toString()),
                        sessionId = sessionId,
                        timestamp = truncNow,
                        schemaVersion = 1,
                        causationId = null,
                        correlationId = null,
                    ),
                    payload = ContextTruncatedEvent(
                        sessionId = sessionId,
                        turnId = userTurnId,
                        entriesDropped = contextPack.compressionMetadata.entriesDropped,
                        truncatedLayers = contextPack.compressionMetadata.truncatedLayers.map { it.name },
                        timestampMs = truncNow.toEpochMilliseconds(),
                    ),
                ),
            )
        }

        val provider = inferenceRouter.route(effectiveStageId, setOf(ModelCapability.General))
        val inferenceRequest = InferenceRequest(
            requestId = InferenceRequestId(UUID.randomUUID().toString()),
            sessionId = sessionId,
            stageId = effectiveStageId,
            contextPack = contextPack,
            generationConfig = config.generationConfig,
            responseFormat = ResponseFormat.Text,
        )
        val inferenceResponse = provider.infer(inferenceRequest)
        val content = inferenceResponse.text

        // Emit ROUTER turn event
        emitChatTurn(
            sessionId, content, ChatTurnRole.ROUTER,
            latencyMs = inferenceResponse.latencyMs,
            tokensUsed = inferenceResponse.tokensUsed,
        )

        if (mode == ChatMode.STEERING) {
            val validationError = validateSteering?.invoke(content)
            if (validationError == null) {
                emitSteeringNote(sessionId, content, effectiveStageId)
            }
        }

        return RouterResponse(content = content, steeringEmitted = (mode == ChatMode.STEERING))
    }

    override suspend fun narrate(sessionId: SessionId, trigger: NarrationTrigger) {
        val state = routerRepository.getRouterState(sessionId)
        val effectiveStageId = state.currentStageId ?: StageId.NONE

        val contextPack = routerContextBuilder.buildNarrationContext(state, trigger, config.tokenBudget)

        val provider = inferenceRouter.route(effectiveStageId, setOf(ModelCapability.General))
        val inferenceRequest = InferenceRequest(
            requestId = InferenceRequestId(UUID.randomUUID().toString()),
            sessionId = sessionId,
            stageId = effectiveStageId,
            contextPack = contextPack,
            generationConfig = config.narrationGenerationConfig,
            responseFormat = ResponseFormat.Text,
        )
        val inferenceResponse = provider.infer(inferenceRequest)
        val content = inferenceResponse.text

        if (content.isBlank()) return

        val narrationId = UUID.randomUUID().toString()
        val now = Clock.System.now()
        eventStore.append(
            NewEvent(
                metadata = EventMetadata(
                    eventId = EventId(UUID.randomUUID().toString()),
                    sessionId = sessionId,
                    timestamp = now,
                    schemaVersion = 1,
                    causationId = null,
                    correlationId = null,
                ),
                payload = RouterNarrationEvent(
                    sessionId = sessionId,
                    narrationId = narrationId,
                    trigger = trigger.kind,
                    stageId = effectiveStageId.takeIf { it != StageId.NONE }?.value,
                    content = content,
                    latencyMs = inferenceResponse.latencyMs,
                    tokensUsed = inferenceResponse.tokensUsed,
                    timestampMs = now.toEpochMilliseconds(),
                ),
            ),
        )
    }

    private suspend fun emitChatTurn(
        sessionId: SessionId,
        content: String,
        role: ChatTurnRole,
        latencyMs: Long? = null,
        tokensUsed: TokenUsage? = null,
    ): String {
        val turnId = UUID.randomUUID().toString()
        val now = Clock.System.now()
        eventStore.append(
            NewEvent(
                metadata = EventMetadata(
                    eventId = EventId(UUID.randomUUID().toString()),
                    sessionId = sessionId,
                    timestamp = now,
                    schemaVersion = 1,
                    causationId = null,
                    correlationId = null,
                ),
                payload = ChatTurnEvent(
                    sessionId = sessionId,
                    turnId = turnId,
                    role = role,
                    content = content,
                    timestampMs = now.toEpochMilliseconds(),
                    latencyMs = latencyMs,
                    tokensUsed = tokensUsed,
                ),
            ),
        )

        runCatching {
            val vector = embedder.embed(content)
            l3MemoryStore.store(
                L3MemoryEntry(
                    id = turnId,
                    sessionId = sessionId,
                    turnId = turnId,
                    text = content,
                    vector = vector,
                    timestampMs = now.toEpochMilliseconds(),
                )
            )
        }.also { result ->
            result.exceptionOrNull()?.let { e ->
                if (e is CancellationException) throw e
                log.warn("L3 write failed for turn {}: {}", turnId, e.message)
            }
        }

        return turnId
    }

    private suspend fun emitSteeringNote(sessionId: SessionId, content: String, effectiveStageId: StageId) {
        val now = Clock.System.now()
        eventStore.append(
            NewEvent(
                metadata = EventMetadata(
                    eventId = EventId(UUID.randomUUID().toString()),
                    sessionId = sessionId,
                    timestamp = now,
                    schemaVersion = 1,
                    causationId = null,
                    correlationId = null,
                ),
                payload = SteeringNoteAddedEvent(
                    sessionId = sessionId,
                    content = content,
                    stageId = effectiveStageId.takeIf { it != StageId.NONE },
                ),
            ),
        )
    }
}

enum class ChatMode {
    CHAT,
    STEERING,
}
