package com.tokilabs.proto_socket

import com.tokilabs.proto_socket.packets.TestData
import kotlinx.coroutines.CompletableDeferred
import kotlinx.coroutines.delay
import kotlinx.coroutines.runBlocking
import java.io.IOException
import java.nio.file.Paths
import kotlin.test.Test
import kotlin.test.assertEquals

class TlsWsTest {
    @Test
    fun testWssSendReceive() = runBlocking {
        val (serverContext, clientContext) = createTlsContexts()
        val port = freePort()
        val received = CompletableDeferred<TestData>()
        val listenerReady = CompletableDeferred<Unit>()
        val server = WsServer("127.0.0.1", port, "/", sslContext = serverContext) { conn ->
            WsClient.forServer(conn, 0, 0, testParserMap())
        }
        server.onClientConnected = { client ->
            addListenerTyped<TestData>(client.communicator) { received.complete(it) }
            listenerReady.complete(Unit)
        }
        var client: WsClient? = null
        server.start()
        try {
            client = dialWssWithRetry("127.0.0.1", port, "/", clientContext)

            listenerReady.await()
            client.send(TestData.newBuilder().setIndex(55).setMessage("hello wss").build())

            assertEquals(55, received.await().index)
        } finally {
            closeClientAndServer(client, server)
        }
    }

    @Test
    fun testWssRequestResponse() = runBlocking {
        val (serverContext, clientContext) = createTlsContexts()
        val port = freePort()
        val server = WsServer("127.0.0.1", port, "/", sslContext = serverContext) { conn ->
            WsClient.forServer(conn, 0, 0, testParserMap())
        }
        server.onClientConnected = { client ->
            addRequestListenerTyped<TestData, TestData>(client.communicator) { req ->
                TestData.newBuilder()
                    .setIndex(req.index * 2)
                    .setMessage("wss echo: ${req.message}")
                    .build()
            }
        }
        var client: WsClient? = null
        server.start()
        try {
            client = dialWssWithRetry("127.0.0.1", port, "/", clientContext)

            val res = sendRequestTyped<TestData, TestData>(
                client.communicator,
                TestData.newBuilder().setIndex(11).setMessage("wss req").build(),
                timeoutMs = 2_000,
            )

            assertEquals(22, res.index)
            assertEquals("wss echo: wss req", res.message)
        } finally {
            closeClientAndServer(client, server)
        }
    }

    private fun createTlsContexts(): Pair<javax.net.ssl.SSLContext, javax.net.ssl.SSLContext> {
        val loader = Thread.currentThread().contextClassLoader
        val certPath = Paths.get(loader.getResource("server.crt")!!.toURI()).toString()
        val keyPath = Paths.get(loader.getResource("server.key")!!.toURI()).toString()
        return createTestSslContexts(certPath, keyPath)
    }

    private suspend fun dialWssWithRetry(
        host: String,
        port: Int,
        path: String,
        sslContext: javax.net.ssl.SSLContext,
    ): WsClient {
        var lastError: IOException? = null
        repeat(3) { attempt ->
            try {
                return dialWss(host, port, path, sslContext, 0, 0, testParserMap())
            } catch (ex: IOException) {
                lastError = ex
                if (attempt < 2) delay(100)
            }
        }
        throw lastError ?: IOException("wss dial failed")
    }

    private suspend fun closeClientAndServer(client: WsClient?, server: WsServer) {
        try {
            client?.close()
            if (client != null) {
                waitForCondition(message = "wss client did not disconnect") {
                    server.clients().isEmpty()
                }
            }
        } finally {
            server.stop()
        }
    }
}
