package com.correx.core.talkie.l3

import kotlinx.coroutines.sync.Mutex
import kotlinx.coroutines.sync.withLock
import kotlin.math.sqrt

/**
 * In-memory L3 memory store for testing. Not production-grade.
 * Uses cosine similarity for vector matching, thread-safe via Mutex.
 */
class InMemoryL3MemoryStore : L3MemoryStore {
    private val entries = mutableListOf<L3MemoryEntry>()
    private val mutex = Mutex()

    override suspend fun store(entry: L3MemoryEntry) {
        mutex.withLock {
            entries.add(entry)
        }
    }

    override suspend fun query(query: L3Query): List<L3Hit> {
        return mutex.withLock {
            entries
                .map { entry ->
                    val score = cosineSimilarity(query.vector, entry.vector)
                    L3Hit(entry, score)
                }
                .filter { hit ->
                    query.sessionIdFilter == null || hit.entry.sessionId == query.sessionIdFilter
                }
                .sortedByDescending { it.score }
                .take(query.k)
        }
    }

    override suspend fun existsByTurnIdPrefix(prefix: String): Boolean =
        mutex.withLock { entries.any { it.turnId.startsWith(prefix) } }

    override suspend fun close() {
        // No-op for in-memory store
    }

    private fun cosineSimilarity(a: FloatArray, b: FloatArray): Float {
        require(a.size == b.size) { "Vector sizes must match" }
        val dotProduct = a.indices.sumOf { (a[it] * b[it]).toDouble() }
        val normA = sqrt(a.indices.sumOf { (a[it] * a[it]).toDouble() })
        val normB = sqrt(b.indices.sumOf { (b[it] * b[it]).toDouble() })
        if (normA == 0.0 || normB == 0.0) return 0f
        return (dotProduct / (normA * normB)).toFloat()
    }
}
