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 kotlinx.serialization.encodeToString
import kotlinx.serialization.json.Json
import com.correx.infrastructure.inference.commons.GpuStatus
import com.correx.infrastructure.inference.commons.ManagedInferenceRouter
import com.correx.infrastructure.inference.commons.ModelDescriptor
import com.correx.infrastructure.inference.commons.ResidencyMode
import com.correx.infrastructure.inference.commons.ResourceProbe
import com.correx.infrastructure.inference.commons.ResourceSnapshot
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.delay
import kotlinx.coroutines.withTimeout
import org.junit.jupiter.api.Test
import org.junit.jupiter.api.io.TempDir
import java.nio.file.Path
import kotlinx.coroutines.runBlocking
import kotlin.test.assertEquals
import kotlin.test.assertNotNull
import kotlin.test.assertTrue
import io.ktor.server.testing.testApplication

/**
 * CI test: exercises the correx-managed model lifecycle WebSocket protocol against a
 * real Ktor [testApplication] server. No actual llama-server binary is spawned — the
 * [DefaultModelManager] uses the system "sleep" binary + a MockEngine health client
 * (mirrors the seam used by DefaultModelManagerTest).
 *
 * Assertions:
 *  1. Initial snapshot contains model.list with [model-a, model-b], current=model-a.
 *  2. Initial snapshot contains resource.status with the fake probe values.
 *  3. SwapModel("model-b") → model.changed arrives with modelId=model-b, loaded=true.
 *  4. ClearModelPin does not produce an error frame.
 */
class ModelLifecycleWiringTest {

    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)

    @Test
    fun `lifecycle protocol over websocket`(@TempDir tempDir: Path): Unit {
        val eventStore = buildInMemoryEventStore()
        val httpClient = buildMockHttpClient()

        val descriptorA = ModelDescriptor(
            modelId = "model-a",
            modelPath = "/dev/null",
            residencyMode = ResidencyMode.PERSISTENT,
            contextSize = 512,
        )
        val descriptorB = ModelDescriptor(
            modelId = "model-b",
            modelPath = "/dev/null",
            residencyMode = ResidencyMode.PERSISTENT,
            contextSize = 512,
        )

        val modelManager = buildTestModelManager(
            llamaBin = "sleep",
            host = "127.0.0.1",
            port = 19888,
            healthTimeoutMs = 5_000L,
            eventStore = eventStore,
            httpClient = httpClient,
        )

        // Pre-load model-a so that currentModelId() returns "model-a" in the initial snapshot,
        // mirroring Main's runBlocking { modelManager.load(descriptor) } call.
        runBlocking { modelManager.load(descriptorA) }

        val descriptors = listOf(descriptorA, descriptorB)
        val modelSwapper = ManagedInferenceRouter(modelManager, descriptors, defaultModelId = "model-a")

        val fakeProbe = ResourceProbe {
            ResourceSnapshot(
                gpu = GpuStatus(memoryUsedMb = 2048L, memoryTotalMb = 8192L, utilizationPct = 42),
                processRssBytes = 512L * 1024L * 1024L,
            )
        }

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

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

            client.webSocket("/stream") {
                // Collect frames until we have both model.list and resource.status
                var modelList: ServerMessage.ModelList? = null
                var resourceStatus: ServerMessage.ResourceStatus? = null

                // Pass 1: collect initial snapshot messages
                withTimeout(5_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")
                assertEquals(listOf("model-a", "model-b"), ml.models)
                assertEquals("model-a", ml.current)

                val rs = assertNotNull(resourceStatus, "resource.status must arrive in initial snapshot")
                assertEquals(2048L, rs.gpuMemoryUsedMb)
                assertEquals(8192L, rs.gpuMemoryTotalMb)
                assertEquals(42, rs.gpuUtilizationPct)
                assertNotNull(rs.processRssMb)

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

                // Wait for model.changed with modelId=model-b, loaded=true
                var modelChanged: ServerMessage.ModelChanged? = null
                withTimeout(8_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: ${msg.message}")
                            else -> Unit
                        }
                    }
                }

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

                // Send ClearModelPin — the server responds with nothing (just clears the pin);
                // verify no ProtocolError arrives within a brief window.
                val clearMsg = encodeClientMessage(ClientMessage.ClearModelPin)
                send(Frame.Text(clearMsg))

                // Give the server a moment to process ClearModelPin; if it sends back a
                // ProtocolError it would arrive almost immediately (well before 2500ms resource push).
                delay(100)
                val pending = incoming.tryReceive()
                if (pending.isSuccess) {
                    val f = pending.getOrThrow()
                    if (f is Frame.Text) {
                        val msg = decodeServerMessage(f.readText())
                        assertTrue(msg !is ServerMessage.ProtocolError, "ClearModelPin produced error: $msg")
                    }
                }
                // If no frame arrived — which is also fine (no error sent).
            }
        }
        } finally {
            runBlocking {
                runCatching { modelManager.unload(modelSwapper.currentModelId() ?: "model-b") }
            }
        }
    }
}
