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.io.BufferedReader
import java.nio.file.Paths
import kotlin.io.path.createDirectories
import kotlin.io.path.exists

/**
 * 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
    // One reader owned by the store for the process's lifetime. A fresh bufferedReader() per request
    // would strand read-ahead bytes in the discarded reader's buffer, desyncing the stream.
    private var reader: BufferedReader? = 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 ->
                // Persist vectors before exit so the next process can `load` them (only save() puts
                // the quantized index on disk — init alone loses everything between runs).
                config.persistPath?.let { runCatching { sendRequestLocked(SidecarRequest(op = "save", path = it.toString())) } }
                runCatching { sendRequestLocked(SidecarRequest(op = "shutdown")) }
                try {
                    withTimeout(SHUTDOWN_TIMEOUT_MS) {
                        p.waitFor()
                    }
                } catch (e: Exception) {
                    if (e !is kotlinx.coroutines.CancellationException) {
                        p.destroyForcibly()
                    } else {
                        throw e
                    }
                }
                reader = null
                process = null
            }
        }
    }

    private suspend fun sendRequest(request: SidecarRequest): SidecarResponse =
        mutex.withLock { sendRequestLocked(request) }

    /** Caller must hold [mutex]. */
    private fun sendRequestLocked(request: SidecarRequest): SidecarResponse {
        ensureProcessStarted()
        val process = process ?: throw RuntimeException("Process failed to start")
        val reader = reader ?: throw RuntimeException("Sidecar reader not initialized")

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

        // A blocking readLine can't be cancelled by a coroutine timeout, so a slow/late response
        // would land in the NEXT request's read and desync the stream permanently. We don't wrap it
        // in withTimeout at all: instead the null/failure path below tears the process down so the
        // next request restarts from a clean, empty stream rather than reading a stale line.
        val responseLine = runCatching { reader.readLine() }.getOrNull()
        if (responseLine == null) {
            killProcess()
            throw RuntimeException("No response from sidecar (process reset)")
        }
        return json.decodeFromString(SidecarResponse.serializer(), responseLine)
    }

    private fun killProcess() {
        runCatching { process?.destroyForcibly() }
        reader = null
        process = null
    }

    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
        reader = started.inputStream.bufferedReader()

        // Reload a persisted index if one exists, else create a fresh one. Either way the sidecar
        // rejects add/search until initialized; without this every store/query fails "not initialized".
        val bootstrap = config.persistPath
            ?.takeIf { it.exists() }
            ?.let { SidecarRequest(op = "load", path = it.toString()) }
            ?: SidecarRequest(op = "init", dim = config.dim, bitWidth = config.bitWidth)

        started.outputStream.write((json.encodeToString(SidecarRequest.serializer(), bootstrap) + "\n").toByteArray())
        started.outputStream.flush()
        val response = reader!!.readLine()?.let { json.decodeFromString(SidecarResponse.serializer(), it) }
        if (response?.ok != true) {
            throw RuntimeException("TurboVec ${bootstrap.op} failed: ${response?.error ?: "no response"}")
        }
    }

    private companion object {
        const val SHUTDOWN_TIMEOUT_MS = 2000L
    }
}
