package com.tokilabs.proto_socket

import com.tokilabs.proto_socket.packets.TestData
import kotlinx.coroutines.CompletableDeferred
import kotlinx.coroutines.async
import kotlinx.coroutines.awaitAll
import kotlinx.coroutines.runBlocking
import kotlin.test.Test
import kotlin.test.assertEquals
import kotlin.test.assertFalse

class WsTest {
    @Test
    fun testWsSendReceive() = runBlocking {
        val port = freePort()
        val received = CompletableDeferred<TestData>()
        val listenerReady = CompletableDeferred<Unit>()
        val server = WsServer("127.0.0.1", port, "/") { conn ->
            WsClient.forServer(conn, 0, 0, testParserMap())
        }
        server.onClientConnected = { client ->
            addListenerTyped<TestData>(client.communicator) { received.complete(it) }
            listenerReady.complete(Unit)
        }
        server.start()
        val client = DialWs("127.0.0.1", port, "/", 0, 0, testParserMap())

        listenerReady.await()
        client.send(TestData.newBuilder().setIndex(42).setMessage("hello proto-socket ws").build())

        assertEquals(42, received.await().index)
        client.close()
        server.stop()
    }

    @Test
    fun testWsRequestResponse() = runBlocking {
        val port = freePort()
        val server = WsServer("127.0.0.1", port, "/") { conn ->
            WsClient.forServer(conn, 0, 0, testParserMap())
        }
        server.onClientConnected = { client ->
            addRequestListenerTyped<TestData, TestData>(client.communicator) { req ->
                TestData.newBuilder()
                    .setIndex(req.index * 2)
                    .setMessage("echo: ${req.message}")
                    .build()
            }
        }
        server.start()
        val client = DialWs("127.0.0.1", port, "/", 0, 0, testParserMap())

        val res = sendRequestTyped<TestData, TestData>(
            client.communicator,
            TestData.newBuilder().setIndex(21).setMessage("hello ws").build(),
            timeoutMs = 2_000,
        )

        assertEquals(42, res.index)
        assertEquals("echo: hello ws", res.message)
        client.close()
        server.stop()
    }

    @Test
    fun testWsLargePayloadRequestResponse() = runBlocking {
        val port = freePort()
        val handlerReady = CompletableDeferred<Unit>()
        val server = WsServer("127.0.0.1", port, "/") { conn ->
            WsClient.forServer(conn, 0, 0, testParserMap())
        }
        server.onClientConnected = { client ->
            addRequestListenerTyped<TestData, TestData>(client.communicator) { req ->
                TestData.newBuilder()
                    .setIndex(req.index + 1)
                    .setMessage(req.message)
                    .build()
            }
            handlerReady.complete(Unit)
        }
        server.start()
        val client = DialWs("127.0.0.1", port, "/", 0, 0, testParserMap())
        val payload = "kotlin-ws-large-payload-".repeat(3_000)

        try {
            handlerReady.await()
            val res = sendRequestTyped<TestData, TestData>(
                client.communicator,
                TestData.newBuilder().setIndex(71).setMessage(payload).build(),
                timeoutMs = 5_000,
            )

            assertEquals(72, res.index)
            assertEquals(payload, res.message)
        } finally {
            client.close()
            server.stop()
        }
    }

    @Test
    fun testWsBroadcast() = runBlocking {
        val port = freePort()
        val server = WsServer("127.0.0.1", port, "/") { conn ->
            WsClient.forServer(conn, 0, 0, testParserMap())
        }
        server.start()
        val client1 = DialWs("127.0.0.1", port, "/", 0, 0, testParserMap())
        val client2 = DialWs("127.0.0.1", port, "/", 0, 0, testParserMap())
        val received1 = CompletableDeferred<TestData>()
        val received2 = CompletableDeferred<TestData>()
        addListenerTyped<TestData>(client1.communicator) { received1.complete(it) }
        addListenerTyped<TestData>(client2.communicator) { received2.complete(it) }

        waitForCondition(message = "clients did not connect") { server.clients().size == 2 }
        server.broadcast(TestData.newBuilder().setIndex(9).setMessage("ws broadcast").build())

        assertEquals("ws broadcast", received1.await().message)
        assertEquals("ws broadcast", received2.await().message)
        client1.close()
        client2.close()
        server.stop()
    }

    @Test
    fun testWsConcurrentRequests() = runBlocking {
        val port = freePort()
        val handlerReady = CompletableDeferred<Unit>()
        val server = WsServer("127.0.0.1", port, "/") { conn ->
            WsClient.forServer(conn, 0, 0, testParserMap())
        }
        server.onClientConnected = { client ->
            addRequestListenerTyped<TestData, TestData>(client.communicator) { req ->
                TestData.newBuilder()
                    .setIndex(req.index * 2)
                    .setMessage("echo: ${req.message}")
                    .build()
            }
            handlerReady.complete(Unit)
        }
        server.start()
        val client = DialWs("127.0.0.1", port, "/", 0, 0, testParserMap())

        handlerReady.await()
        val deferred = (0 until 5).map { i ->
            async {
                val index = 30 + i
                val message = "request $i"
                val res = sendRequestTyped<TestData, TestData>(
                    client.communicator,
                    TestData.newBuilder().setIndex(index).setMessage(message).build(),
                    timeoutMs = 2_000,
                )
                assertEquals(index * 2, res.index)
                assertEquals("echo: $message", res.message)
            }
        }
        try {
            deferred.awaitAll()
        } finally {
            client.close()
            server.stop()
        }
    }

    @Test
    fun testWsServerStopDisconnectsClients() = runBlocking {
        val port = freePort()
        val server = WsServer("127.0.0.1", port, "/") { conn ->
            WsClient.forServer(conn, 0, 0, testParserMap())
        }
        server.start()
        val client = DialWs("127.0.0.1", port, "/", 0, 0, testParserMap())
        val disconnected = CompletableDeferred<Unit>()
        client.addDisconnectListener { disconnected.complete(Unit) }

        waitForCondition(message = "client did not connect") { server.clients().size == 1 }
        server.stop()

        disconnected.await()
        assertFalse(client.communicator.isAlive())
    }
}
