import com.correx.core.events.events.ChatTurnEvent
import com.correx.core.events.events.ChatTurnRole
import com.correx.core.events.events.EventMetadata
import com.correx.core.events.events.EventPayload
import com.correx.core.events.events.L3MemoryRetrievedEvent
import com.correx.core.events.events.NewEvent
import com.correx.core.events.events.StoredEvent
import com.correx.core.events.stores.EventStore
import com.correx.core.events.types.EventId
import com.correx.core.events.types.SessionId
import com.correx.core.talkie.l3.InMemoryL3MemoryStore
import com.correx.core.talkie.l3.L3Hit
import com.correx.core.talkie.l3.L3MemoryEntry
import com.correx.core.talkie.l3.L3MetadataRehydrator
import com.correx.core.talkie.l3.L3Query
import com.correx.core.talkie.l3.RehydratableL3MemoryStore
import kotlinx.coroutines.flow.Flow
import kotlinx.coroutines.runBlocking
import kotlinx.datetime.Instant
import org.junit.jupiter.api.Assertions.assertEquals
import org.junit.jupiter.api.Assertions.assertTrue
import org.junit.jupiter.api.Test

class L3MetadataRehydratorTest {

    private class CapturingStore : RehydratableL3MemoryStore {
        val loaded = mutableListOf<L3MemoryEntry>()
        override suspend fun rehydrateMetadata(entries: List<L3MemoryEntry>) {
            loaded += entries
        }
        override suspend fun store(entry: L3MemoryEntry) = Unit
        override suspend fun query(query: L3Query): List<L3Hit> = emptyList()
        override suspend fun existsByTurnIdPrefix(prefix: String): Boolean = false
        override suspend fun close() = Unit
    }

    private class FakeEventStore(private val stored: List<StoredEvent>) : EventStore {
        override fun allEvents(): Sequence<StoredEvent> = stored.asSequence()
        override fun read(sessionId: SessionId): List<StoredEvent> =
            stored.filter { it.metadata.sessionId == sessionId }
        override fun readFrom(sessionId: SessionId, fromSequence: Long): List<StoredEvent> = read(sessionId)
        override fun lastSequence(sessionId: SessionId): Long? = null
        override fun allSessionIds(): Set<SessionId> = stored.map { it.metadata.sessionId }.toSet()
        override suspend fun append(event: NewEvent): StoredEvent = throw UnsupportedOperationException()
        override suspend fun appendAll(events: List<NewEvent>): List<StoredEvent> = throw UnsupportedOperationException()
        override fun subscribe(sessionId: SessionId): Flow<StoredEvent> = throw UnsupportedOperationException()
        override fun subscribeAll(): Flow<StoredEvent> = throw UnsupportedOperationException()
        override suspend fun lastGlobalSequence(): Long = 0
    }

    private var seq = 0L

    private fun event(sessionId: String, payload: EventPayload): StoredEvent {
        seq += 1
        return StoredEvent(
            metadata = EventMetadata(
                eventId = EventId("e$seq"),
                sessionId = SessionId(sessionId),
                timestamp = Instant.fromEpochMilliseconds(seq),
                schemaVersion = 1,
                causationId = null,
                correlationId = null,
            ),
            sequence = seq,
            sessionSequence = seq,
            payload = payload,
        )
    }

    private fun chatTurn(sessionId: String, turnId: String, text: String) =
        event(sessionId, ChatTurnEvent(SessionId(sessionId), turnId, ChatTurnRole.USER, text, timestampMs = 42L))

    @Test
    fun `rehydrates chat turns across sessions and ignores non-chat events`(): Unit = runBlocking {
        val store = FakeEventStore(
            listOf(
                chatTurn("s1", "t1", "hello"),
                event("s1", L3MemoryRetrievedEvent(SessionId("s1"), "t1", emptyList(), timestampMs = 1L)),
                chatTurn("s2", "t2", "world"),
            ),
        )
        val target = CapturingStore()

        val count = L3MetadataRehydrator(store).rehydrate(target)

        assertEquals(2, count)
        assertEquals(2, target.loaded.size)
        val byTurn = target.loaded.associateBy { it.turnId }
        assertEquals(SessionId("s1"), byTurn.getValue("t1").sessionId)
        assertEquals("hello", byTurn.getValue("t1").text)
        assertEquals("t1", byTurn.getValue("t1").id)
        assertEquals(SessionId("s2"), byTurn.getValue("t2").sessionId)
        assertTrue(byTurn.getValue("t2").vector.isEmpty(), "rehydrated entry holds no vector")
    }

    @Test
    fun `non-rehydratable store is skipped`(): Unit = runBlocking {
        val store = FakeEventStore(listOf(chatTurn("s1", "t1", "hello")))

        val count = L3MetadataRehydrator(store).rehydrate(InMemoryL3MemoryStore())

        assertEquals(0, count)
    }
}
