proto-socket/kotlin/crosstest/kotlin_python.kt

@file:JvmName("KotlinPythonKt")

package com.tokilabs.toki_socket.crosstest

import com.tokilabs.toki_socket.TcpClient
import com.tokilabs.toki_socket.TcpServer
import com.tokilabs.toki_socket.ParserMap
import com.tokilabs.toki_socket.WsClient
import com.tokilabs.toki_socket.WsServer
import com.tokilabs.toki_socket.addListenerTyped
import com.tokilabs.toki_socket.addRequestListenerTyped
import com.tokilabs.toki_socket.typeNameOf
import com.tokilabs.toki_socket.packets.TestData
import kotlinx.coroutines.CompletableDeferred
import kotlinx.coroutines.Dispatchers
import kotlinx.coroutines.coroutineScope
import kotlinx.coroutines.delay
import kotlinx.coroutines.launch
import kotlinx.coroutines.runBlocking
import kotlinx.coroutines.withContext
import kotlinx.coroutines.withTimeoutOrNull
import java.io.File
import java.util.concurrent.TimeUnit
import kotlin.system.exitProcess

private const val HOST = "127.0.0.1"
private const val TCP_PORT = 29394
private const val WS_PORT = 29396
private const val WS_PATH = "/"
private const val PROCESS_TIMEOUT_MS = 20_000L
private const val SERVER_OBSERVATION_MS = 200L

fun main() = runBlocking {
    println("INFO typeName kotlin=${typeNameOf<TestData>()}")
    try {
        runTcpSendPush()
        delay(150)
        runTcpRequests()
        delay(150)
        runWsSendPush()
        delay(150)
        runWsRequests()
        println("PASS all kotlin-server/python-client crosstests passed")
    } catch (error: Throwable) {
        System.err.println("FAIL crosstest error=${error.message ?: error}")
        exitProcess(1)
    }
}

private suspend fun runTcpSendPush() = coroutineScope {
    val received = CompletableDeferred<Boolean>()
    val server = TcpServer(HOST, TCP_PORT) { socket ->
        TcpClient.fromSocket(socket, 0, 0, parserMap())
    }
    server.onClientConnected = { client ->
        addListenerTyped<TestData>(client.communicator) { data ->
            println("SERVER_RECEIVED index=${data.index} message=${data.message}")
            val valid = data.index == 101 && data.message == "fire from python client"
            received.complete(valid)
            if (valid) {
                launch {
                    client.send(
                        TestData.newBuilder()
                            .setIndex(200)
                            .setMessage("push from kotlin server")
                            .build(),
                    )
                }
            }
        }
    }
    withServer(server::start, server::stop) {
        runPythonClient("tcp", TCP_PORT, "send-push", setOf("1", "2"))
        val ok = withTimeoutOrNull(SERVER_OBSERVATION_MS) { received.await() } ?: false
        check(ok) { "TCP send-push server did not receive expected data" }
    }
}

private suspend fun runTcpRequests() {
    val server = TcpServer(HOST, TCP_PORT) { socket ->
        TcpClient.fromSocket(socket, 0, 0, parserMap())
    }
    server.onClientConnected = { client ->
        addRequestListenerTyped<TestData, TestData>(client.communicator) { req ->
            TestData.newBuilder()
                .setIndex(req.index * 2)
                .setMessage("echo: ${req.message}")
                .build()
        }
    }
    withServer(server::start, server::stop) {
        runPythonClient("tcp", TCP_PORT, "requests", setOf("3", "4"))
    }
}

private suspend fun runWsSendPush() = coroutineScope {
    val received = CompletableDeferred<Boolean>()
    val server = WsServer(HOST, WS_PORT, WS_PATH) { conn ->
        WsClient.forServer(conn, 0, 0, parserMap())
    }
    server.onClientConnected = { client ->
        addListenerTyped<TestData>(client.communicator) { data ->
            println("SERVER_RECEIVED index=${data.index} message=${data.message}")
            val valid = data.index == 101 && data.message == "fire from python client"
            received.complete(valid)
            if (valid) {
                launch {
                    client.send(
                        TestData.newBuilder()
                            .setIndex(200)
                            .setMessage("push from kotlin server")
                            .build(),
                    )
                }
            }
        }
    }
    withServer(server::start, server::stop) {
        runPythonClient("ws", WS_PORT, "send-push", setOf("1", "2"))
        val ok = withTimeoutOrNull(SERVER_OBSERVATION_MS) { received.await() } ?: false
        check(ok) { "WS send-push server did not receive expected data" }
    }
}

private suspend fun runWsRequests() {
    val server = WsServer(HOST, WS_PORT, WS_PATH) { conn ->
        WsClient.forServer(conn, 0, 0, parserMap())
    }
    server.onClientConnected = { client ->
        addRequestListenerTyped<TestData, TestData>(client.communicator) { req ->
            TestData.newBuilder()
                .setIndex(req.index * 2)
                .setMessage("echo: ${req.message}")
                .build()
        }
    }
    withServer(server::start, server::stop) {
        runPythonClient("ws", WS_PORT, "requests", setOf("3", "4"))
    }
}

private suspend fun withServer(
    start: () -> Unit,
    stop: () -> Unit,
    body: suspend () -> Unit,
) {
    start()
    try {
        body()
    } finally {
        stop()
    }
}

private suspend fun runPythonClient(
    mode: String,
    port: Int,
    phase: String,
    expectedScenarios: Set<String>,
) = withContext(Dispatchers.IO) {
    val process = ProcessBuilder(
        "python3",
        "crosstest/kotlin_python_client.py",
        "--mode=$mode",
        "--port=$port",
        "--phase=$phase",
    )
        .directory(pythonDir())
        .redirectErrorStream(false)
        .start()

    val stdoutLines = mutableListOf<String>()
    val stdoutThread = Thread {
        process.inputStream.bufferedReader().forEachLine { line ->
            println(line)
            if (line.startsWith("PASS ") || line.startsWith("FAIL ")) {
                stdoutLines.add(line)
            }
        }
    }
    val stderrThread = Thread {
        process.errorStream.bufferedReader().forEachLine { line ->
            System.err.println(line)
        }
    }
    stdoutThread.start()
    stderrThread.start()

    val finished = process.waitFor(PROCESS_TIMEOUT_MS, TimeUnit.MILLISECONDS)
    if (!finished) {
        process.destroyForcibly()
    }
    stdoutThread.join()
    stderrThread.join()

    validateResultLines(
        "python-client $mode/$phase",
        if (finished) process.exitValue() else -1,
        stdoutLines,
        expectedScenarios,
    )
}

private fun validateResultLines(
    label: String,
    exitCode: Int,
    lines: List<String>,
    expectedScenarios: Set<String>,
) {
    val failed = lines.filter { it.startsWith("FAIL ") }
    val passed = lines
        .filter { it.startsWith("PASS ") }
        .mapNotNull { Regex("""scenario=([^ ]+)""").find(it)?.groupValues?.get(1) }
        .toSet()
    val missing = expectedScenarios - passed
    check(exitCode == 0 && failed.isEmpty() && missing.isEmpty()) {
        "$label failed exitCode=$exitCode failed=$failed missing=$missing"
    }
}

private fun parserMap(): ParserMap =
    mapOf(typeNameOf<TestData>() to { TestData.parseFrom(it) })

private fun pythonDir(): File {
    var dir = File(System.getProperty("user.dir")).absoluteFile
    while (dir.parentFile != null) {
        val candidate = File(dir, "../python").canonicalFile
        if (File(candidate, "pyproject.toml").isFile) return candidate
        val direct = File(dir, "python").canonicalFile
        if (File(direct, "pyproject.toml").isFile) return direct
        dir = dir.parentFile
    }
    error("cannot resolve python directory")
}