package com.correx.core.talkie.l3

import com.correx.core.events.types.SessionId
import com.correx.core.utils.TypeId
import kotlinx.coroutines.runBlocking
import org.junit.jupiter.api.Test
import org.junit.jupiter.api.Assertions.assertEquals

class InMemoryL3MemoryStoreTest {
    @Test
    fun `stores and queries entries by cosine similarity`(): Unit = runBlocking {
        val store = InMemoryL3MemoryStore()
        val sessionId = TypeId("session-1")

        val entry1 = L3MemoryEntry(
            id = "entry-1",
            sessionId = sessionId,
            turnId = "turn-1",
            text = "hello world",
            vector = floatArrayOf(1f, 0f, 0f),
            timestampMs = 100L
        )

        val entry2 = L3MemoryEntry(
            id = "entry-2",
            sessionId = sessionId,
            turnId = "turn-2",
            text = "goodbye world",
            vector = floatArrayOf(0.9f, 0.1f, 0f),
            timestampMs = 200L
        )

        val entry3 = L3MemoryEntry(
            id = "entry-3",
            sessionId = sessionId,
            turnId = "turn-3",
            text = "far away",
            vector = floatArrayOf(0f, 0f, 1f),
            timestampMs = 300L
        )

        store.store(entry1)
        store.store(entry2)
        store.store(entry3)

        val query = L3Query(vector = floatArrayOf(1f, 0f, 0f), k = 2)
        val hits = store.query(query)

        assertEquals(2, hits.size, "Should return k=2 results")
        assertEquals("entry-1", hits[0].entry.id, "Closest match should be entry-1")
        assertEquals("entry-2", hits[1].entry.id, "Second match should be entry-2")
    }

    @Test
    fun `existsByTurnIdPrefix returns true when entry turnId starts with prefix`(): Unit = runBlocking {
        val store = InMemoryL3MemoryStore()
        val sessionId = TypeId("session-1")
        store.store(
            L3MemoryEntry(
                id = "e1", sessionId = sessionId, turnId = "repomap:/repo:git:abc123",
                text = "foo", vector = floatArrayOf(1f), timestampMs = 1L
            )
        )
        assert(store.existsByTurnIdPrefix("repomap:/repo:git:abc123")) { "should match exact" }
        assert(store.existsByTurnIdPrefix("repomap:/repo:git:")) { "should match prefix" }
    }

    @Test
    fun `existsByTurnIdPrefix returns false when no entry matches`(): Unit = runBlocking {
        val store = InMemoryL3MemoryStore()
        val sessionId = TypeId("session-1")
        store.store(
            L3MemoryEntry(
                id = "e1", sessionId = sessionId, turnId = "project:/repo",
                text = "foo", vector = floatArrayOf(1f), timestampMs = 1L
            )
        )
        assert(!store.existsByTurnIdPrefix("repomap:/repo:git:")) { "should not match different prefix" }
    }

    @Test
    fun `respects sessionIdFilter`(): Unit = runBlocking {
        val store = InMemoryL3MemoryStore()
        val session1 = TypeId("session-1")
        val session2 = TypeId("session-2")

        val entry1 = L3MemoryEntry(
            id = "entry-1",
            sessionId = session1,
            turnId = "turn-1",
            text = "session 1",
            vector = floatArrayOf(1f, 0f),
            timestampMs = 100L
        )

        val entry2 = L3MemoryEntry(
            id = "entry-2",
            sessionId = session2,
            turnId = "turn-1",
            text = "session 2",
            vector = floatArrayOf(1f, 0f),
            timestampMs = 200L
        )

        store.store(entry1)
        store.store(entry2)

        val query = L3Query(vector = floatArrayOf(1f, 0f), k = 10, sessionIdFilter = session1)
        val hits = store.query(query)

        assertEquals(1, hits.size, "Should only return entry from session1")
        assertEquals("entry-1", hits[0].entry.id)
    }
}
