package com.tokilabs.proto_socket

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

class TcpTest {
    @Test
    fun testTcpSendReceive() = runBlocking {
        val port = freePort()
        val received = CompletableDeferred<TestData>()
        val listenerReady = CompletableDeferred<Unit>()
        val server = TcpServer("127.0.0.1", port) { socket ->
            TcpClient.fromSocket(socket, 0, 0, testParserMap())
        }
        server.onClientConnected = { client ->
            addListenerTyped<TestData>(client.communicator) { received.complete(it) }
            listenerReady.complete(Unit)
        }
        server.start()
        val client = dialTcpWithRetry("127.0.0.1", port)

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

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

    @Test
    fun testTcpRequestResponse() = runBlocking {
        val port = freePort()
        val handlerReady = CompletableDeferred<Unit>()
        val server = TcpServer("127.0.0.1", port) { socket ->
            TcpClient.fromSocket(socket, 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 = dialTcpWithRetry("127.0.0.1", port)

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

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

    @Test
    fun testTcpBroadcast() = runBlocking {
        val port = freePort()
        val server = TcpServer("127.0.0.1", port) { socket ->
            TcpClient.fromSocket(socket, 0, 0, testParserMap())
        }
        server.start()
        val client = dialTcpWithRetry("127.0.0.1", port)
        val received = CompletableDeferred<TestData>()
        addListenerTyped<TestData>(client.communicator) { received.complete(it) }

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

        assertEquals("broadcast", received.await().message)
        client.close()
        server.stop()
    }

    @Test
    fun testTcpServerStopDisconnectsClients() = runBlocking {
        val port = freePort()
        val server = TcpServer("127.0.0.1", port) { socket ->
            TcpClient.fromSocket(socket, 0, 0, testParserMap())
        }
        server.start()
        val client = dialTcpWithRetry("127.0.0.1", port)
        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())
    }

    @Test
    fun testTcpConcurrentRequests() = runBlocking {
        val port = freePort()
        val handlerReady = CompletableDeferred<Unit>()
        val server = TcpServer("127.0.0.1", port) { socket ->
            TcpClient.fromSocket(socket, 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 = dialTcpWithRetry("127.0.0.1", port)

        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 testTcpClientCloseIdempotent() = runBlocking {
        val serverSocket = java.net.ServerSocket(0)
        val accept = async(Dispatchers.IO) { serverSocket.accept() }
        val socket = Socket("127.0.0.1", serverSocket.localPort)
        val peer = accept.await()
        val client = TcpClient.fromSocket(socket, 0, 0, testParserMap())

        repeat(3) { client.close() }

        assertFalse(client.communicator.isAlive())
        peer.close()
        serverSocket.close()
    }

    private suspend fun dialTcpWithRetry(
        host: String,
        port: Int,
    ): TcpClient {
        var lastError: java.io.IOException? = null
        repeat(5) { attempt ->
            try {
                return DialTcp(host, port, 0, 0, testParserMap())
            } catch (ex: java.io.IOException) {
                lastError = ex
                if (attempt < 4) kotlinx.coroutines.delay(100)
            }
        }
        throw lastError ?: java.io.IOException("tcp dial failed")
    }
}
