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.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) {
            scope.launch {
                eventStore.subscribe(sessionId).collect { event ->
                    processEvent(sessionId, event)
                }
            }
        }
    }

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