package com.correx.infrastructure.artifactscas.compact

import com.correx.core.events.events.EventMetadata
import com.correx.core.events.events.InferenceCompletedEvent
import com.correx.core.events.events.InferenceStartedEvent
import com.correx.core.events.events.NewEvent
import com.correx.core.events.types.ArtifactId
import com.correx.core.events.types.EventId
import com.correx.core.events.types.SessionId
import com.correx.core.inference.TokenUsage
import com.correx.core.utils.TypeId
import com.correx.infrastructure.artifactscas.CasArtifactStore
import com.correx.infrastructure.artifactscas.config.CasConfig
import com.correx.infrastructure.artifactscas.index.SqliteArtifactIndex
import com.correx.infrastructure.persistence.InMemoryEventStore
import kotlinx.coroutines.runBlocking
import kotlinx.datetime.Clock
import org.junit.jupiter.api.Assertions.assertArrayEquals
import org.junit.jupiter.api.Assertions.assertEquals
import org.junit.jupiter.api.Assertions.assertFalse
import org.junit.jupiter.api.Assertions.assertNull
import org.junit.jupiter.api.Assertions.assertTrue
import org.junit.jupiter.api.Test
import org.junit.jupiter.api.io.TempDir
import java.nio.file.Files
import java.nio.file.Path
import java.util.*
import kotlin.io.path.listDirectoryEntries

class CompactorTest {

    private val sessionId: SessionId = TypeId("session-1")

    private fun newStore(dir: Path, maxSegmentBytes: Long = 64L * 1024 * 1024): CasArtifactStore {
        val index = SqliteArtifactIndex(dir.resolve("index.sqlite"))
        return CasArtifactStore(CasConfig(dir, maxSegmentBytes), index)
    }

    private suspend fun emitStarted(store: InMemoryEventStore, prompt: ArtifactId) {
        store.append(
            NewEvent(
                metadata = EventMetadata(
                    eventId = EventId(UUID.randomUUID().toString()),
                    sessionId = sessionId,
                    timestamp = Clock.System.now(),
                    schemaVersion = 1,
                    causationId = null,
                    correlationId = null,
                ),
                payload = InferenceStartedEvent(
                    requestId = TypeId(UUID.randomUUID().toString()),
                    sessionId = sessionId,
                    stageId = TypeId("stage"),
                    providerId = TypeId("prov"),
                    promptArtifactId = prompt,
                ),
            )
        )
    }

    private suspend fun emitCompleted(store: InMemoryEventStore, response: ArtifactId) {
        store.append(
            NewEvent(
                metadata = EventMetadata(
                    eventId = EventId(UUID.randomUUID().toString()),
                    sessionId = sessionId,
                    timestamp = Clock.System.now(),
                    schemaVersion = 1,
                    causationId = null,
                    correlationId = null,
                ),
                payload = InferenceCompletedEvent(
                    requestId = TypeId(UUID.randomUUID().toString()),
                    sessionId = sessionId,
                    stageId = TypeId("stage"),
                    providerId = TypeId("prov"),
                    tokensUsed = TokenUsage(0, 0),
                    latencyMs = 0L,
                    responseArtifactId = response,
                ),
            )
        )
    }

    @Test
    fun `no compaction when all artifacts live`(@TempDir dir: Path): Unit = runBlocking {
        newStore(dir, maxSegmentBytes = 256).use { store ->
            val eventStore = InMemoryEventStore()
            val a = store.put("aaaa".toByteArray())
            val b = store.put("bbbb".toByteArray())
            val c = store.put("cccc".toByteArray())
            emitCompleted(eventStore, a)
            emitCompleted(eventStore, b)
            emitCompleted(eventStore, c)
            val report = store.compact(eventStore)
            assertTrue(report.compactedSegmentIds.isEmpty(), "report: $report")
        }
    }

    @Test
    fun `compacts a mostly-dead sealed segment`(@TempDir dir: Path): Unit = runBlocking {
        newStore(dir, maxSegmentBytes = 64).use { store ->
            val eventStore = InMemoryEventStore()
            val live = store.put(ByteArray(40) { 1 })
            val dead = store.put(ByteArray(40) { 2 })
            store.put(ByteArray(40) { 3 })
            emitCompleted(eventStore, live)

            val segsBefore = dir.resolve("segments").listDirectoryEntries("*.seg")
                .map { it.fileName.toString() }.toSet()
            assertTrue(segsBefore.size >= 3, "expected rotation, got $segsBefore")

            val report = store.compact(eventStore)
            assertFalse(report.compactedSegmentIds.isEmpty(), "expected sealed compactions, got $report")

            assertArrayEquals(ByteArray(40) { 1 }, store.get(live))
            assertNull(store.get(dead))
        }
    }

    @Test
    fun `fully dead sealed segment is deleted with no new segment written`(@TempDir dir: Path): Unit = runBlocking {
        val sealedName: String
        val sealedBytes: Long
        val sealedSegId: Long
        newStore(dir, maxSegmentBytes = 64).use { store ->
            val eventStore = InMemoryEventStore()
            val dead = store.put(ByteArray(40) { 1 })
            store.put(ByteArray(40) { 2 })

            val segsDir = dir.resolve("segments")
            val segsBefore = segsDir.listDirectoryEntries("*.seg")
                .associate { it.fileName.toString() to Files.size(it) }
            assertEquals(2, segsBefore.size, "expected 2 segments before compaction, got $segsBefore")
            sealedName = segsBefore.keys.min()
            sealedBytes = segsBefore.getValue(sealedName)
            sealedSegId = sealedName.removeSuffix(".seg").toLong()

            val report = store.compact(eventStore)

            assertEquals(listOf(sealedSegId), report.compactedSegmentIds, "report: $report")
            assertNull(report.newSegmentId, "no new segment should be written, got $report")
            assertEquals(sealedBytes, report.bytesReclaimed, "report: $report")

            val segsAfter = segsDir.listDirectoryEntries("*.seg")
                .map { it.fileName.toString() }.toSet()
            assertFalse(sealedName in segsAfter, "sealed segment file should be deleted, got $segsAfter")
            assertEquals(1, segsAfter.size, "only active segment should remain, got $segsAfter")

            assertNull(store.get(dead), "dead artifact must not be retrievable")
        }
        SqliteArtifactIndex(dir.resolve("index.sqlite")).use { reopened ->
            assertTrue(
                reopened.entriesIn(sealedSegId).isEmpty(),
                "index entries for sealed segment must be removed",
            )
        }
    }

    @Test
    fun `never touches the active segment`(@TempDir dir: Path): Unit = runBlocking {
        newStore(dir, maxSegmentBytes = 64L * 1024 * 1024).use { store ->
            val eventStore = InMemoryEventStore()
            store.put("x".toByteArray())
            store.put("y".toByteArray())
            store.put("z".toByteArray())

            val segsDir = dir.resolve("segments")
            val before = segsDir.listDirectoryEntries("*.seg")
                .associate { it.fileName.toString() to Files.size(it) }

            val report = store.compact(eventStore)
            assertTrue(report.compactedSegmentIds.isEmpty(), "active must not be compacted, got $report")

            val after = segsDir.listDirectoryEntries("*.seg")
                .associate { it.fileName.toString() to Files.size(it) }
            assertEquals(before, after)
        }
    }

    @Test
    fun `liveness scanner counts both prompt and response ids`(@TempDir dir: Path): Unit = runBlocking {
        val eventStore = InMemoryEventStore()
        val id1: ArtifactId
        val id2: ArtifactId
        newStore(dir).use { store ->
            id1 = store.put("p1".toByteArray())
            id2 = store.put("p2".toByteArray())
            emitStarted(eventStore, id1)
            emitCompleted(eventStore, id2)
        }
        val index = SqliteArtifactIndex(dir.resolve("index.sqlite"))
        index.use {
            val scanner = LivenessScanner(eventStore, index)
            val live = scanner.scan().liveIds
            assertEquals(setOf(id1, id2), live)
        }
    }
}
