package com.correx.infrastructure.artifactscas.compact

import com.correx.core.events.events.InferenceCompletedEvent
import com.correx.core.events.events.InferenceStartedEvent
import com.correx.core.events.stores.EventStore
import com.correx.core.events.types.ArtifactId
import com.correx.infrastructure.artifactscas.index.ArtifactIndex
import kotlinx.coroutines.Dispatchers
import kotlinx.coroutines.withContext
import java.nio.file.Files
import java.nio.file.Path
import kotlin.io.path.listDirectoryEntries

class LivenessScanner(
    private val eventStore: EventStore,
    private val index: ArtifactIndex,
) {
    data class SegmentStats(val liveBytes: Long, val totalBytes: Long) {
        val liveRatio: Double get() = if (totalBytes == 0L) 1.0 else liveBytes.toDouble() / totalBytes
    }

    data class Liveness(
        val liveIds: Set<ArtifactId>,
        val perSegment: Map<Long, SegmentStats>,
    )

    suspend fun scan(): Liveness = scan(segmentsDir = null)

    suspend fun scan(segmentsDir: Path?): Liveness {
        val liveIds = collectLiveIds()
        val perSegmentLive = mutableMapOf<Long, Long>()
        for (id in liveIds) {
            val hash = id.value.hexToBytes()
            val loc = index.lookup(hash) ?: continue
            perSegmentLive.merge(loc.segmentId, loc.length.toLong()) { a, b -> a + b }
        }
        val perSegment = if (segmentsDir != null && Files.exists(segmentsDir)) {
            segmentsDir.listDirectoryEntries("*.seg")
                .mapNotNull { p ->
                    val id = runCatching { p.fileName.toString().removeSuffix(".seg").toLong() }.getOrNull()
                        ?: return@mapNotNull null
                    val total = runCatching { Files.size(p) }.getOrDefault(0L)
                    val live = perSegmentLive[id] ?: 0L
                    id to SegmentStats(live, total)
                }.toMap()
        } else {
            perSegmentLive.mapValues { (_, live) -> SegmentStats(live, live) }
        }
        return Liveness(liveIds, perSegment)
    }

    private suspend fun collectLiveIds(): Set<ArtifactId> = withContext(Dispatchers.IO) {
        val ids = mutableSetOf<ArtifactId>()
        for (event in eventStore.allEvents()) {
            when (val p = event.payload) {
                is InferenceStartedEvent -> ids.add(p.promptArtifactId)
                is InferenceCompletedEvent -> ids.add(p.responseArtifactId)
                else -> Unit
            }
        }
        ids
    }

}

private const val HEX_RADIX = 16
private const val HEX_CHARS_PER_BYTE = 2
private const val HEX_NIBBLE_BITS = 4

internal fun String.hexToBytes(): ByteArray {
    require(length % HEX_CHARS_PER_BYTE == 0) { "hex string must have even length" }
    return ByteArray(length / HEX_CHARS_PER_BYTE) { i ->
        val hi = Character.digit(this[HEX_CHARS_PER_BYTE * i], HEX_RADIX)
        val lo = Character.digit(this[HEX_CHARS_PER_BYTE * i + 1], HEX_RADIX)
        ((hi shl HEX_NIBBLE_BITS) + lo).toByte()
    }
}
