package com.correx.infrastructure.router.turbovec

import com.correx.core.talkie.l3.L3Hit
import com.correx.core.talkie.l3.L3MemoryEntry
import com.correx.core.talkie.l3.L3Query
import com.correx.core.talkie.l3.RehydratableL3MemoryStore
import kotlinx.coroutines.sync.Mutex
import kotlinx.coroutines.sync.withLock
import kotlinx.coroutines.withTimeout
import kotlinx.serialization.json.Json
import java.nio.file.Paths
import kotlin.io.path.createDirectories

/**
 * Adapter to TurboVec sidecar process for cross-session vector-based memory.
 * Maintains in-memory metadata (sessionId, turnId, text, timestamp) since turbovec
 * stores only vectors and IDs.
 *
 * The metadata map is volatile, but the sidecar vectors persist (via `persistPath`);
 * on restart [rehydrateMetadata] rebuilds the map from the `ChatTurnEvent` log so
 * queries resolve their hits again (invariant #1: the event log is authoritative).
 */
class TurboVecL3MemoryStore(private val config: TurboVecSidecarConfig) : RehydratableL3MemoryStore {
    private val json = Json { ignoreUnknownKeys = true }
    private val mutex = Mutex()
    private var process: Process? = null
    private val metadata = java.util.concurrent.ConcurrentHashMap<String, L3MemoryEntry>()

    override suspend fun rehydrateMetadata(entries: List<L3MemoryEntry>) {
        entries.forEach { metadata.putIfAbsent(it.id, it) }
    }

    override suspend fun store(entry: L3MemoryEntry) {
        val request = SidecarRequest(
            op = "add",
            id = entry.id,
            vector = entry.vector.toList()
        )
        val response = sendRequest(request)
        if (!response.ok) {
            throw RuntimeException("Failed to store entry: ${response.error}")
        }
        metadata[entry.id] = entry
    }

    override suspend fun existsByTurnIdPrefix(prefix: String): Boolean =
        metadata.values.any { it.turnId.startsWith(prefix) }

    override suspend fun query(query: L3Query): List<L3Hit> {
        val request = SidecarRequest(
            op = "search",
            vector = query.vector.toList(),
            k = query.k
        )
        val response = sendRequest(request)
        if (!response.ok) {
            throw RuntimeException("Failed to query: ${response.error}")
        }

        return (response.hits ?: emptyList())
            .mapNotNull { hit ->
                metadata[hit.id]?.let { entry ->
                    L3Hit(entry, hit.score)
                }
            }
            .filter { hit ->
                query.sessionIdFilter == null || hit.entry.sessionId == query.sessionIdFilter
            }
            .sortedByDescending { it.score }
    }

    override suspend fun close() {
        mutex.withLock {
            process?.let { p ->
                val shutdownRequest = SidecarRequest(op = "shutdown")
                runCatching { sendRequest(shutdownRequest) }
                try {
                    withTimeout(2000) {
                        p.waitFor()
                    }
                } catch (e: Exception) {
                    if (e !is kotlinx.coroutines.CancellationException) {
                        p.destroyForcibly()
                    } else {
                        throw e
                    }
                }
            }
        }
    }

    private suspend fun sendRequest(request: SidecarRequest): SidecarResponse {
        mutex.withLock {
            ensureProcessStarted()
            val process = process ?: throw RuntimeException("Process failed to start")

            val requestLine = json.encodeToString(SidecarRequest.serializer(), request) + "\n"
            process.outputStream.write(requestLine.toByteArray())
            process.outputStream.flush()

            val responseLine = runCatching {
                withTimeout(config.requestTimeoutMs) {
                    process.inputStream.bufferedReader().readLine()
                }
            }.getOrNull() ?: throw RuntimeException("No response from sidecar")

            return json.decodeFromString(SidecarResponse.serializer(), responseLine)
        }
    }

    private fun ensureProcessStarted() {
        if (process != null && process!!.isAlive) {
            return
        }

        val logDir = Paths.get(System.getProperty("user.home"), ".local", "share", "correx", "logs")
        logDir.createDirectories()
        val logFile = logDir.resolve("turbovec-sidecar.log").toFile()

        val pb = ProcessBuilder(config.pythonExecutable, config.scriptPath.toString())
        pb.redirectError(logFile)

        val started = pb.start()
        process = started

        // The sidecar rejects add/search until the index is created. Send init once on startup;
        // without this every store/query fails with "Index not initialized".
        val initLine = json.encodeToString(
            SidecarRequest.serializer(),
            SidecarRequest(op = "init", dim = config.dim, bitWidth = config.bitWidth)
        ) + "\n"
        started.outputStream.write(initLine.toByteArray())
        started.outputStream.flush()
        val response = started.inputStream.bufferedReader().readLine()
            ?.let { json.decodeFromString(SidecarResponse.serializer(), it) }
        if (response?.ok != true) {
            throw RuntimeException("TurboVec init failed: ${response?.error ?: "no response"}")
        }
    }
}
