package com.correx.apps.server.memory

import com.correx.core.events.events.RepoKnowledgeHit
import com.correx.core.events.types.SessionId
import com.correx.core.inference.Embedder
import com.correx.core.kernel.orchestration.RepoKnowledgeRetriever
import com.correx.core.router.l3.L3MemoryStore
import com.correx.core.router.l3.L3Query

private const val RETRIEVAL_OVERSAMPLE_FACTOR = 4

class L3RepoKnowledgeRetriever(
    private val embedder: Embedder,
    private val l3MemoryStore: L3MemoryStore,
    private val repoRoot: String,
) : RepoKnowledgeRetriever {
    override suspend fun retrieve(sessionId: SessionId, query: String, k: Int): List<RepoKnowledgeHit> {
        val vector = embedder.embed(query)
        return l3MemoryStore.query(L3Query(vector = vector, k = k * RETRIEVAL_OVERSAMPLE_FACTOR))
            .filter { it.entry.turnId.startsWith("repomap:$repoRoot:") }
            .take(k)
            .map { RepoKnowledgeHit(path = it.entry.text.substringBefore(":"), text = it.entry.text, score = it.score) }
    }
}
