package com.correx.apps.server.lifecycle

import com.correx.apps.server.configureServer
import com.correx.apps.server.protocol.ClientMessage
import com.correx.apps.server.protocol.ServerMessage
import com.correx.infrastructure.inference.commons.ManagedInferenceRouter
import com.correx.infrastructure.inference.commons.ModelDescriptor
import com.correx.infrastructure.inference.commons.NvidiaResourceProbe
import com.correx.infrastructure.inference.commons.ResidencyMode
import com.correx.infrastructure.inference.commons.UnavailableProbe
import io.ktor.client.engine.cio.CIO
import io.ktor.client.plugins.websocket.WebSockets
import io.ktor.client.plugins.websocket.webSocket
import io.ktor.websocket.Frame
import io.ktor.websocket.readText
import kotlinx.coroutines.runBlocking
import kotlinx.coroutines.withTimeout
import kotlinx.serialization.encodeToString
import kotlinx.serialization.json.Json
import org.junit.jupiter.api.AfterEach
import org.junit.jupiter.api.Assumptions
import org.junit.jupiter.api.Test
import org.junit.jupiter.api.condition.EnabledIfEnvironmentVariable
import org.junit.jupiter.api.io.TempDir
import java.nio.file.Path
import kotlin.test.assertEquals
import kotlin.test.assertNotNull
import kotlin.test.assertTrue
import io.ktor.server.testing.testApplication

/**
 * Live integration test: spawns a real llama-server process and exercises the correx-managed
 * model lifecycle WebSocket protocol end-to-end.
 *
 * Gated on the [CORREX_TEST_LLAMA_BIN] environment variable; skipped automatically in CI.
 *
 * **Run command:**
 * ```
 * CORREX_TEST_LLAMA_BIN=/path/to/llama-server \
 * CORREX_TEST_MODEL=/path/to/model-a.gguf \
 * CORREX_TEST_MODEL_2=/path/to/model-b.gguf \
 * ./gradlew :apps:server:test --tests '*ModelLifecycleLiveTest*' --rerun-tasks
 * ```
 *
 * [CORREX_TEST_MODEL_2] is optional. When absent, the swap assertion is skipped via
 * [Assumptions.assumeTrue] and only the load + telemetry assertions run.
 */
@EnabledIfEnvironmentVariable(named = "CORREX_TEST_LLAMA_BIN", matches = ".+")
class ModelLifecycleLiveTest {

    private val llamaBin: String = System.getenv("CORREX_TEST_LLAMA_BIN")
    private val modelAPath: String = System.getenv("CORREX_TEST_MODEL") ?: ""
    private val modelBPath: String? = System.getenv("CORREX_TEST_MODEL_2")
    private val modelPort: Int = System.getenv("CORREX_TEST_MODEL_PORT")?.toIntOrNull() ?: 10000

    private var modelManager: com.correx.infrastructure.inference.llama.cpp.DefaultModelManager? = null
    private var currentModelId: String? = null

    private val protocolJson = Json {
        classDiscriminator = "type"
        ignoreUnknownKeys = true
    }

    private fun decodeServerMessage(text: String): ServerMessage =
        protocolJson.decodeFromString(text)

    private fun encodeClientMessage(msg: ClientMessage): String =
        protocolJson.encodeToString(msg)

    @AfterEach
    fun teardown() {
        val mgr = modelManager ?: return
        val id = currentModelId ?: return
        runBlocking {
            runCatching { mgr.unload(id) }
        }
    }

    @Test
    fun `lifecycle protocol over websocket — live llama-server`(@TempDir tempDir: Path) {
        val eventStore = buildInMemoryEventStore()
        val httpClient = io.ktor.client.HttpClient(CIO)

        val descriptorA = ModelDescriptor(
            modelId = "live-model-a",
            modelPath = modelAPath,
            residencyMode = ResidencyMode.PERSISTENT,
            contextSize = 512,
        )

        val descriptors = if (modelBPath != null) {
            listOf(
                descriptorA,
                ModelDescriptor(
                    modelId = "live-model-b",
                    modelPath = modelBPath,
                    residencyMode = ResidencyMode.PERSISTENT,
                    contextSize = 512,
                ),
            )
        } else {
            listOf(descriptorA)
        }

        val mgr = buildTestModelManager(
            llamaBin = llamaBin,
            host = "127.0.0.1",
            port = modelPort,
            healthTimeoutMs = 120_000L,
            eventStore = eventStore,
            httpClient = httpClient,
        )
        modelManager = mgr
        currentModelId = "live-model-a"

        val modelSwapper = ManagedInferenceRouter(mgr, descriptors, defaultModelId = "live-model-a")

        val resourceProbe = if (NvidiaResourceProbe.isAvailable()) {
            NvidiaResourceProbe(pidSupplier = { mgr.currentPid() })
        } else {
            UnavailableProbe
        }

        val module = buildTestServerModule(
            eventStore = eventStore,
            modelSwapper = modelSwapper,
            resourceProbe = resourceProbe,
            tempDir = tempDir,
        )

        runBlocking { mgr.load(descriptorA) }

        testApplication {
            application { configureServer(module) }
            val client = createClient { install(WebSockets) }

            client.webSocket("/stream") {
                var modelList: ServerMessage.ModelList? = null
                var resourceStatus: ServerMessage.ResourceStatus? = null

                withTimeout(10_000L) {
                    while (modelList == null || resourceStatus == null) {
                        val frame = incoming.receive()
                        if (frame !is Frame.Text) continue
                        val raw = frame.readText()
                        when (val msg = decodeServerMessage(raw)) {
                            is ServerMessage.ModelList -> modelList = msg
                            is ServerMessage.ResourceStatus -> resourceStatus = msg
                            else -> Unit
                        }
                    }
                }

                val ml = assertNotNull(modelList, "model.list must arrive in initial snapshot")
                assertTrue(ml.models.contains("live-model-a"), "model-a must be listed")
                assertEquals("live-model-a", ml.current)

                assertNotNull(resourceStatus, "resource.status must arrive in initial snapshot")

                // Swap assertion: only when a second model is configured
                Assumptions.assumeTrue(
                    modelBPath != null,
                    "CORREX_TEST_MODEL_2 not set — skipping swap assertion",
                )

                currentModelId = "live-model-b"
                val swapMsg = encodeClientMessage(ClientMessage.SwapModel("live-model-b"))
                send(Frame.Text(swapMsg))

                var modelChanged: ServerMessage.ModelChanged? = null
                withTimeout(120_000L) {
                    while (modelChanged == null) {
                        val frame = incoming.receive()
                        if (frame !is Frame.Text) continue
                        val raw = frame.readText()
                        when (val msg = decodeServerMessage(raw)) {
                            is ServerMessage.ModelChanged -> modelChanged = msg
                            is ServerMessage.ProtocolError ->
                                error("Unexpected protocol error during swap: ${msg.message}")
                            else -> Unit
                        }
                    }
                }

                val mc = assertNotNull(modelChanged, "model.changed must arrive after SwapModel")
                assertEquals("live-model-b", mc.modelId)
                assertTrue(mc.loaded, "loaded must be true after swap")

                val clearMsg = encodeClientMessage(ClientMessage.ClearModelPin)
                send(Frame.Text(clearMsg))
            }
        }
    }
}
