package com.correx.infrastructure.persistence.artifact

import com.correx.core.artifacts.ArtifactProjector
import com.correx.core.artifacts.ArtifactReducer
import com.correx.core.artifacts.ArtifactState
import com.correx.core.artifacts.repository.ArtifactRepository
import com.correx.core.events.events.ArtifactCreatedEvent
import com.correx.core.events.events.ArtifactValidatedEvent
import com.correx.core.events.events.ArtifactValidatingEvent
import com.correx.core.events.events.StoredEvent
import com.correx.core.events.stores.EventStore
import com.correx.core.events.types.ArtifactId
import com.correx.core.events.types.SessionId
import com.correx.core.events.types.StageId
import kotlinx.coroutines.CompletableDeferred
import kotlinx.coroutines.CoroutineScope
import kotlinx.coroutines.Dispatchers
import kotlinx.coroutines.Job
import kotlinx.coroutines.SupervisorJob
import kotlinx.coroutines.launch
import java.util.concurrent.*

class LiveArtifactRepository(
    private val eventStore: EventStore,
    reducer: ArtifactReducer,
    private val scope: CoroutineScope = CoroutineScope(Dispatchers.IO + SupervisorJob()),
) : ArtifactRepository {

    private val artifactCache = ConcurrentHashMap<SessionId, ConcurrentHashMap<ArtifactId, ArtifactState>>()
    private val stageIndex = ConcurrentHashMap<SessionId, ConcurrentHashMap<ArtifactId, StageId>>()
    private val subscriptions = ConcurrentHashMap<SessionId, Job>()
    private val projector = ArtifactProjector(reducer)

    override fun getBySession(sessionId: SessionId): Map<ArtifactId, ArtifactState> {
        ensureSubscribed(sessionId)
        return artifactCache[sessionId]?.toMap() ?: emptyMap()
    }

    override fun getByStage(sessionId: SessionId, stageId: StageId): Map<ArtifactId, ArtifactState> {
        ensureSubscribed(sessionId)
        val cache = artifactCache[sessionId] ?: emptyMap()
        val index = stageIndex[sessionId] ?: emptyMap()
        return cache.filter { (artifactId, _) -> index[artifactId] == stageId }
    }

    override fun getById(sessionId: SessionId, artifactId: ArtifactId): ArtifactState? {
        ensureSubscribed(sessionId)
        return artifactCache[sessionId]?.get(artifactId)
    }

    override fun getByIds(sessionId: SessionId, artifactIds: Set<ArtifactId>): Map<ArtifactId, ArtifactState> {
        ensureSubscribed(sessionId)
        val cache = artifactCache[sessionId] ?: return emptyMap()
        return artifactIds.mapNotNull { id -> cache[id]?.let { id to it } }.toMap()
    }

    private fun ensureSubscribed(sessionId: SessionId) {
        subscriptions.computeIfAbsent(sessionId) {
            // Rebuild from persisted history so a cold (post-restart) or first mid-session consumer
            // sees prior artifact state immediately, not empty (invariant #1). The reducer's phase
            // transitions are strictly ordered, so application must stay single-threaded and in
            // sequence order — hence the gate below rather than a concurrent replay.
            val rebuildWatermark = CompletableDeferred<Long>()
            val job = scope.launch {
                // Start collecting immediately so events appended during the rebuild are captured (flow
                // back-pressure buffers them), but hold each until the rebuild's high-water mark is
                // known, then drop anything already folded by the rebuild (sequence <= watermark).
                // Live delivery is in append order, so post-watermark events fold in order too.
                eventStore.subscribe(sessionId).collect { event ->
                    if (event.sequence > rebuildWatermark.await()) processEvent(sessionId, event)
                }
            }
            var watermark = 0L
            eventStore.read(sessionId).forEach { event ->
                watermark = maxOf(watermark, event.sequence)
                processEvent(sessionId, event)
            }
            rebuildWatermark.complete(watermark)
            job
        }
    }

    private fun processEvent(sessionId: SessionId, event: StoredEvent) {
        val artifactId = extractArtifactId(event) ?: return
        val sessionArtifacts = artifactCache.computeIfAbsent(sessionId) { ConcurrentHashMap() }
        val currentState = sessionArtifacts[artifactId] ?: projector.initial()
        val newState = projector.apply(currentState, event)
        sessionArtifacts[artifactId] = newState

        val payload = event.payload
        if (payload is ArtifactCreatedEvent) {
            val sessionStages = stageIndex.computeIfAbsent(sessionId) { ConcurrentHashMap() }
            sessionStages[artifactId] = payload.stageId
        }
    }

    private fun extractArtifactId(event: StoredEvent): ArtifactId? =
        when (val p = event.payload) {
            is ArtifactCreatedEvent -> p.artifactId
            is ArtifactValidatingEvent -> p.artifactId
            is ArtifactValidatedEvent -> p.artifactId
            else -> null
        }
}
