package com.correx.apps.server.memory

import com.correx.core.events.types.SessionId
import com.correx.core.inference.Embedder
import com.correx.core.talkie.l3.InMemoryL3MemoryStore
import com.correx.core.talkie.l3.L3MemoryEntry
import kotlinx.coroutines.runBlocking
import org.junit.jupiter.api.Assertions.assertEquals
import org.junit.jupiter.api.Test

private class RetrieverConstantEmbedder(override val dimension: Int = 4) : Embedder {
    override suspend fun embed(text: String): FloatArray = FloatArray(dimension) { 0.5f }
}

class L3RepoKnowledgeRetrieverTest {

    private val vec = FloatArray(4) { 0.5f }

    @Test
    fun `retriever for a repoRoot does not match a sibling whose path it prefixes`(): Unit = runBlocking {
        val l3 = InMemoryL3MemoryStore()
        l3.store(L3MemoryEntry("1", SessionId("s"), "repomap:v2:/repo:git:h", "repo/A.kt: Foo", vec, 0L))
        l3.store(L3MemoryEntry("2", SessionId("s"), "repomap:v2:/repo2:git:h", "repo2/B.kt: Bar", vec, 0L))

        val hits = L3RepoKnowledgeRetriever(RetrieverConstantEmbedder(), l3, "/repo")
            .retrieve(SessionId("s"), "anything", 10)

        assertEquals(listOf("repo/A.kt: Foo"), hits.map { it.text }, "must not leak /repo2 entries into /repo")
    }

    @Test
    fun `retriever matches its own repoRoot entries with or without a stateKey suffix`(): Unit = runBlocking {
        val l3 = InMemoryL3MemoryStore()
        l3.store(L3MemoryEntry("1", SessionId("s"), "repomap:v2:/repo:git:h1", "repo/A.kt", vec, 0L))
        l3.store(L3MemoryEntry("2", SessionId("s"), "repomap:v2:/repo:", "repo/B.kt", vec, 0L))

        val hits = L3RepoKnowledgeRetriever(RetrieverConstantEmbedder(), l3, "/repo")
            .retrieve(SessionId("s"), "anything", 10)

        assertEquals(setOf("repo/A.kt", "repo/B.kt"), hits.map { it.text }.toSet())
    }
}
