package com.correx.apps.server.protocol

import com.correx.core.events.types.SessionId
import kotlinx.serialization.json.Json
import org.junit.jupiter.api.Assertions.assertEquals
import org.junit.jupiter.api.Test

class ServerMessageSerializationTest {

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

    @Test
    fun `SessionAnnounced encodes sequence and sessionSequence`() {
        val msg = ServerMessage.SessionAnnounced(
            sessionId = SessionId("sess-1"),
            workflowId = "wf-1",
            sequence = 42,
            sessionSequence = 7,
        )
        val jsonStr = ProtocolSerializer.encodeServerMessage(msg)
        assert(jsonStr.contains("\"type\":\"session.announced\"")) { "expected type=session.announced" }
        assert(jsonStr.contains("\"sessionId\":\"sess-1\"")) { "expected sessionId" }
        assert(jsonStr.contains("\"workflowId\":\"wf-1\"")) { "expected workflowId" }
        assert(jsonStr.contains("\"sequence\":42")) { "expected sequence=42" }
        assert(jsonStr.contains("\"sessionSequence\":7")) { "expected sessionSequence=7" }
    }

    @Test
    fun `ChatTurn encodes role content and cursors`() {
        val msg = ServerMessage.ChatTurn(
            sessionId = SessionId("sess-1"),
            turnId = "turn-9",
            role = "ROUTER",
            content = "hi there",
            sequence = 5,
            sessionSequence = 2,
        )
        val jsonStr = ProtocolSerializer.encodeServerMessage(msg)
        assert(jsonStr.contains("\"type\":\"chat.turn\"")) { "expected type=chat.turn" }
        assert(jsonStr.contains("\"turnId\":\"turn-9\"")) { "expected turnId" }
        assert(jsonStr.contains("\"role\":\"ROUTER\"")) { "expected role" }
        assert(jsonStr.contains("\"content\":\"hi there\"")) { "expected content" }
    }

    @Test
    fun `ChatTurn with metrics encodes latencyMs and totalTokens`() {
        val msg = ServerMessage.ChatTurn(
            sessionId = SessionId("sess-1"),
            turnId = "t9",
            role = "ROUTER",
            content = "hi",
            latencyMs = 3000,
            totalTokens = 943,
            sequence = 7,
            sessionSequence = 2,
        )
        val jsonStr = ProtocolSerializer.encodeServerMessage(msg)
        assert(jsonStr.contains("\"type\":\"chat.turn\"")) { "expected type=chat.turn" }
        assert(jsonStr.contains("\"latencyMs\":3000")) { "expected latencyMs=3000" }
        assert(jsonStr.contains("\"totalTokens\":943")) { "expected totalTokens=943" }
    }

    @Test
    fun `ChatTurn without metrics encodes latencyMs and totalTokens as absent or null`() {
        val msg = ServerMessage.ChatTurn(
            sessionId = SessionId("sess-1"),
            turnId = "t10",
            role = "USER",
            content = "hello",
            sequence = 8,
            sessionSequence = 3,
        )
        val jsonStr = ProtocolSerializer.encodeServerMessage(msg)
        assert(jsonStr.contains("\"type\":\"chat.turn\"")) { "expected type=chat.turn" }
        assert(!jsonStr.contains("\"latencyMs\":3000")) { "latencyMs should not be 3000" }
        assert(!jsonStr.contains("\"totalTokens\":943")) { "totalTokens should not be 943" }
    }

    @Test
    fun `SnapshotComplete encodes with only type field`() {
        val msg = ServerMessage.SnapshotComplete
        val jsonStr = ProtocolSerializer.encodeServerMessage(msg)
        assert(jsonStr.contains("\"type\":\"snapshot_complete\"")) { "expected type=snapshot_complete" }
    }

    @Test
    fun `SessionSnapshot encodes lastSequence and lastSessionSequence`() {
        val msg = ServerMessage.SessionSnapshot(
            sessionId = SessionId("sess-1"),
            workflowId = "wf-abc",
            state = SessionStateDto(
                status = "running",
                currentStageId = "stage-1",
                pauseReason = null,
            ),
            pendingApprovals = emptyList(),
            lastSequence = 100,
            lastSessionSequence = 15,
        )
        val jsonStr = ProtocolSerializer.encodeServerMessage(msg)
        assert(jsonStr.contains("\"type\":\"session_snapshot\"")) { "expected type=session_snapshot" }
        assert(jsonStr.contains("\"sessionId\":\"sess-1\"")) { "expected sessionId" }
        assert(jsonStr.contains("\"lastSequence\":100")) { "expected lastSequence=100" }
        assert(jsonStr.contains("\"lastSessionSequence\":15")) { "expected lastSessionSequence=15" }
    }

    @Test
    fun `Narration encodes type content stageId latencyMs and totalTokens`() {
        val msg = ServerMessage.Narration(
            sessionId = SessionId("sess-1"),
            content = "moving on",
            stageId = "validate",
            latencyMs = 2000,
            totalTokens = 120,
            sequence = 10,
            sessionSequence = 5,
        )
        val jsonStr = ProtocolSerializer.encodeServerMessage(msg)
        assert(jsonStr.contains("\"type\":\"router.narration\"")) { "expected type=router.narration" }
        assert(jsonStr.contains("\"content\":\"moving on\"")) { "expected content" }
        assert(jsonStr.contains("\"latencyMs\":2000")) { "expected latencyMs=2000" }
        assert(jsonStr.contains("\"totalTokens\":120")) { "expected totalTokens=120" }
    }

    @Test
    fun `SessionAnnounced round-trips through JSON`() {
        val original = ServerMessage.SessionAnnounced(
            sessionId = SessionId("sess-42"),
            workflowId = "wf-abc",
            sequence = 200,
            sessionSequence = 33,
        )
        val jsonStr = ProtocolSerializer.encodeServerMessage(original)
        val decoded = json.decodeFromString<ServerMessage.SessionAnnounced>(jsonStr)
        assertEquals(original.sessionId, decoded.sessionId)
        assertEquals(original.workflowId, decoded.workflowId)
        assertEquals(200L, decoded.sequence)
        assertEquals(33L, decoded.sessionSequence)
    }
}
