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 testTcpCloseCompletesWhileReadParked() = runBlocking {
        // read 루프가 soTimeout poll로 readInt에 parked된 상태에서도 close가 socket close에만
        // 의존하지 않고 즉시 완료되는지 확인한다(parallel clients>=128 unskip 회귀 방지).
        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)
        waitForCondition(message = "client did not connect") { server.clients().size == 1 }

        // poll 주기를 한 번 이상 지나 read가 확실히 blocking readInt에 머무는 시점에 close한다.
        kotlinx.coroutines.delay(1_200)
        val elapsed = kotlin.system.measureTimeMillis { client.close() }

        assertFalse(client.communicator.isAlive())
        kotlin.test.assertTrue(elapsed < 2_000, "close blocked for ${elapsed}ms")
        server.stop()
    }

    @Test
    fun testTcpDelayedPartialFrameDoesNotDesync() = runBlocking {
        // frame header/payload를 두 chunk로 나눠 보내되, chunk 사이 지연을 read poll timeout(1000ms)보다
        // 길게 둔다. offset을 보존하지 않으면 chunk 사이 SocketTimeoutException 후 frame 중간 bytes를
        // length로 오해해 desync된다. header split과 payload split을 모두 검증한다.
        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())
        try {
            val first = CompletableDeferred<TestData>()
            val second = CompletableDeferred<TestData>()
            val results = listOf(first, second)
            val nextResult = java.util.concurrent.atomic.AtomicInteger(0)
            addListenerTyped<TestData>(client.communicator) { td ->
                val i = nextResult.getAndIncrement()
                if (i < results.size) results[i].complete(td)
            }

            val out = peer.getOutputStream()

            // 1) payload split: header + payload 앞부분을 보내고, poll timeout을 넘긴 뒤 나머지 payload flush.
            val framePayloadSplit = buildTestDataFrame(7, "payload-split")
            val payloadCut = framePayloadSplit.size - 3
            out.write(framePayloadSplit, 0, payloadCut)
            out.flush()
            kotlinx.coroutines.delay(1_200)
            out.write(framePayloadSplit, payloadCut, framePayloadSplit.size - payloadCut)
            out.flush()
            assertEquals("payload-split", first.await().message)

            // 2) header split: 4-byte length prefix 자체를 둘로 쪼개 poll timeout을 넘겨도 desync되지 않는지 확인.
            val frameHeaderSplit = buildTestDataFrame(9, "header-split")
            out.write(frameHeaderSplit, 0, 2)
            out.flush()
            kotlinx.coroutines.delay(1_200)
            out.write(frameHeaderSplit, 2, frameHeaderSplit.size - 2)
            out.flush()
            assertEquals("header-split", second.await().message)

            kotlin.test.assertTrue(client.communicator.isAlive(), "client disconnected after delayed partial frames")
        } finally {
            client.close()
            peer.close()
            serverSocket.close()
        }
    }

    // 4-byte big-endian length prefix + PacketBase(TestData) envelope. wire 포맷은 Communicator.send와
    // TcpClient.writePacket이 만드는 것과 동일하다.
    private fun buildTestDataFrame(index: Int, message: String): ByteArray {
        val td = TestData.newBuilder().setIndex(index).setMessage(message).build()
        val base = com.tokilabs.proto_socket.packets.PacketBase.newBuilder()
            .setTypeName(typeNameOf<TestData>())
            .setNonce(1)
            .setData(td.toByteString())
            .build()
        val body = base.toByteArray()
        val header = java.nio.ByteBuffer.allocate(4).putInt(body.size).array()
        return header + body
    }

    @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")
    }
}
