package com.correx.infrastructure.artifactscas.compact

import com.correx.infrastructure.artifactscas.config.CasConfig
import com.correx.infrastructure.artifactscas.index.ArtifactIndex
import com.correx.infrastructure.artifactscas.segment.Location
import com.correx.infrastructure.artifactscas.segment.SegmentLayout
import com.correx.infrastructure.artifactscas.segment.SegmentReader
import com.correx.infrastructure.artifactscas.segment.SegmentWriter
import kotlinx.coroutines.Dispatchers
import kotlinx.coroutines.sync.Mutex
import kotlinx.coroutines.sync.withLock
import kotlinx.coroutines.withContext
import java.io.RandomAccessFile
import java.nio.file.Files
import java.nio.file.Path
import kotlin.io.path.listDirectoryEntries

data class CompactionReport(
    val compactedSegmentIds: List<Long>,
    val newSegmentId: Long?,
    val bytesReclaimed: Long,
)

class Compactor(
    config: CasConfig,
    private val index: ArtifactIndex,
    private val liveness: LivenessScanner,
    private val storeMutex: Mutex,
    private val liveRatioThreshold: Double = DEFAULT_LIVE_RATIO_THRESHOLD,
) {
    private val segmentsDir: Path = config.rootDir.resolve("segments")
    private val reader = SegmentReader()

    suspend fun compact(): CompactionReport = withContext(Dispatchers.IO) {
        if (!Files.exists(segmentsDir)) return@withContext empty()
        val segmentIds = listSegmentIds()
        if (segmentIds.isEmpty()) return@withContext empty()
        val activeId = segmentIds.max()
        val scan = liveness.scan(segmentsDir)
        val eligible = selectEligible(scan, activeId)
        if (eligible.isEmpty()) return@withContext empty()

        val liveEntries = readLiveEntries(eligible, scan.liveIds)
        val newSegmentId = activeId + 1
        val (bytesWritten, newLocations) = writeNewSegment(newSegmentId, liveEntries)

        val reclaimedBefore = eligible.sumOf { id -> scan.perSegment[id]?.totalBytes ?: 0L }
        applyAtomicSwap(newLocations, eligible)

        val effectiveNew = if (liveEntries.isEmpty()) {
            runCatching { Files.deleteIfExists(segmentsDir.resolve(SegmentWriter.segmentFileName(newSegmentId))) }
            null
        } else {
            newSegmentId
        }

        CompactionReport(
            compactedSegmentIds = eligible,
            newSegmentId = effectiveNew,
            bytesReclaimed = reclaimedBefore - bytesWritten,
        )
    }

    private fun selectEligible(scan: LivenessScanner.Liveness, activeId: Long): List<Long> =
        scan.perSegment
            .filter { (id, stats) ->
                id != activeId && stats.totalBytes > 0 && stats.liveRatio < liveRatioThreshold
            }
            .keys
            .sorted()

    private suspend fun readLiveEntries(
        eligible: List<Long>,
        liveIds: Set<com.correx.core.events.types.ArtifactId>,
    ): List<LiveEntry> {
        val liveHashHex = liveIds.map { it.value }.toHashSet()
        val out = mutableListOf<LiveEntry>()
        for (segId in eligible) {
            for ((hash, loc) in index.entriesIn(segId)) {
                if (hash.toHex() in liveHashHex) {
                    val bytes = reader.read(loc, segmentsDir)
                    out.add(LiveEntry(hash, bytes))
                }
            }
        }
        return out
    }

    private suspend fun writeNewSegment(
        newSegmentId: Long,
        entries: List<LiveEntry>,
    ): Pair<Long, List<Pair<ByteArray, Location>>> {
        if (entries.isEmpty()) return 0L to emptyList()
        val newPath = segmentsDir.resolve(SegmentWriter.segmentFileName(newSegmentId))
        var bytesWritten = 0L
        val newLocations = mutableListOf<Pair<ByteArray, Location>>()
        withContext(Dispatchers.IO) {
            RandomAccessFile(newPath.toFile(), "rw").use { raf ->
                for (entry in entries) {
                    val header = SegmentLayout.encodeHeader(entry.payload, entry.hash)
                    val offset = raf.filePointer
                    raf.write(header)
                    raf.write(entry.payload)
                    newLocations.add(entry.hash to Location(newSegmentId, offset, entry.payload.size))
                    bytesWritten += header.size + entry.payload.size
                }
                raf.fd.sync()
            }
        }
        return bytesWritten to newLocations
    }

    private suspend fun applyAtomicSwap(newLocations: List<Pair<ByteArray, Location>>, eligible: List<Long>) {
        storeMutex.withLock {
            for ((hash, loc) in newLocations) {
                index.update(hash, loc)
            }
            for (segId in eligible) {
                runCatching {
                    Files.deleteIfExists(segmentsDir.resolve(SegmentWriter.segmentFileName(segId)))
                }
            }
            for (segId in eligible) {
                index.deleteEntriesIn(segId)
            }
            index.flush()
        }
    }

    private fun listSegmentIds(): List<Long> =
        segmentsDir.listDirectoryEntries("*.seg")
            .mapNotNull { p -> runCatching { p.fileName.toString().removeSuffix(".seg").toLong() }.getOrNull() }

    private fun empty(): CompactionReport = CompactionReport(emptyList(), null, 0L)

    private class LiveEntry(
        val hash: ByteArray,
        val payload: ByteArray,
    )

    companion object {
        const val DEFAULT_LIVE_RATIO_THRESHOLD = 0.5
    }
}

private fun ByteArray.toHex(): String = joinToString("") { "%02x".format(it) }
