kiloparsec/src/commonMain/kotlin/net/sergeych/kiloparsec/adapter/websocketClient.kt

/*
 * Copyright (c) 2025-2026. Sergey S. Chernov - All Rights Reserved
 *
 *  You may use, distribute and modify this code under the
 * terms of the private license, which you must obtain from the author
 *
 * To obtain the license, contact the author: https://t.me/real_sergeych or email to
 * real dot sergeych at gmail.
 */

package net.sergeych.kiloparsec.adapter

import io.ktor.client.*
import io.ktor.client.plugins.websocket.*
import io.ktor.http.*
import io.ktor.websocket.*
import kotlinx.coroutines.CancellationException
import kotlinx.coroutines.CompletableDeferred
import kotlinx.coroutines.channels.Channel
import kotlinx.coroutines.channels.ClosedReceiveChannelException
import kotlinx.coroutines.channels.ClosedSendChannelException
import kotlinx.coroutines.launch
import kotlinx.io.IOException
import net.sergeych.crypto2.SigningKey
import net.sergeych.kiloparsec.*
import net.sergeych.mp_logger.*
import net.sergeych.mp_tools.decodeBase64Compact
import net.sergeych.mp_tools.encodeToBase64Compact
import net.sergeych.mp_tools.globalLaunch
import net.sergeych.tools.AtomicCounter

private val counter = AtomicCounter()

/**
 * Shortcut to create websocket client. Use [websocketTransportDevice] with [KiloClient]
 * for fine-grained control.
 */
fun <S> websocketClient(
    path: String,
    clientInterface: KiloInterface<S> = KiloInterface(),
    secretKey: SigningKey? = null,
    useTextFrames: Boolean = false,
    sessionMaker: () -> S = {
        @Suppress("UNCHECKED_CAST")
        Unit as S
    },
): KiloClient<S> {
    return KiloClient(clientInterface, secretKey) {
        KiloConnectionData(websocketTransportDevice(path, useTextFrames), sessionMaker())
    }
}

/**
 * Create kilopaarsec transport over websocket (ws or wss).
 * @param path websocket path (must start with ws:// or wss:// and contain a path part)
 * @param client optional caller-owned client. When omitted, a client with the
 * [WebSockets] plugin is created for this device and closed with it.
 */
fun websocketTransportDevice(
    path: String,
    useTextFrames: Boolean = false,
    client: HttpClient? = null,
): Transport.Device {

    val ownsClient = client == null
    val actualClient = client ?: HttpClient { install(WebSockets) }

    val log = LogTag("WSTD")
    var u = Url(path)
    log.debug { "Creating websocket transport device at $u" }
    if (u.encodedPath.length <= 1) {
        log.debug {"Correcting path as up = ${u.encodedPath}" }
        u = URLBuilder(u).apply {
            encodedPath = "/kp"
        }.build()
    }
    log.debug { "Url to process is $u" }

    val input = Channel<UByteArray>()
    val output = Channel<UByteArray>()
    val closeHandle = CompletableDeferred<Boolean>()
    val readyHandle = CompletableDeferred<Unit>()

    globalLaunch {
        val log = LogTag("KC:${counter.incrementAndGet()}")
        try {
            actualClient.webSocket({
                url.protocol = u.protocol
                url.host = u.host
                url.port = u.port
                url.encodedPath = u.encodedPath
                url.parameters.appendAll(u.parameters)
                log.info { "kiloparsec server URL: $url" }
            }) {
                log.info { "connected to the server" }
//                println("SENDING!!!")
//                send("Helluva")
                readyHandle.complete(Unit)
                launch {
                    try {
                        for (block in output) {
                            if (useTextFrames)
                                send(
                                    Frame.Text(block.asByteArray().encodeToBase64Compact())
                                )
                            else
                                send(block.toByteArray())
                        }
                        log.info { "input is closed, closing the websocket" }
                        if (closeHandle.isActive) closeHandle.complete(true)
                    } catch (_: ClosedSendChannelException) {
                        log.info { "send channel closed" }
                    } catch (_: CancellationException) {
                    } catch (t: Throwable) {
                        log.info { "unexpected exception in websock sender: ${t.stackTraceToString()}" }
                        closeHandle.completeExceptionally(t)
                    }
                    if (closeHandle.isActive) closeHandle.complete(false)
                }
                launch {
                    try {
                        for (f in incoming) {
                            when (f) {
                                is Frame.Binary -> input.send(f.readBytes().toUByteArray())
                                is Frame.Text -> input.send(f.readText().decodeBase64Compact().toUByteArray())
                                else -> log.warning { "ignoring unexpected frame of type ${f.frameType}" }
                            }
                        }
                        if (closeHandle.isActive) closeHandle.complete(true)
                    } catch (_: CancellationException) {
                        if (closeHandle.isActive) closeHandle.complete(false)
                    } catch (_: ClosedReceiveChannelException) {
                        log.warning { "receive channel closed unexpectedly" }
                        if (closeHandle.isActive) closeHandle.complete(false)
                    } catch (t: Throwable) {
                        log.exception { "unexpected error" to t }
                        if (closeHandle.isActive) closeHandle.complete(false)
                    }
                }
                if (!closeHandle.await()) {
                    log.warning { "Client is closing with error" }
                    throw RemoteInterface.ClosedException()
                }
                runCatching { output.close() }
                runCatching { input.close() }
                runCatching { close() }
            }
        } catch (x: IOException) {
            if ("refused" in x.toString()) log.debug { "connection refused" }
            else log.warning { "unexpected IO error $x" }
            runCatching { output.close() }
            runCatching { input.close() }
        } finally {
            if (ownsClient) actualClient.close()
        }
        log.info { "closing connection" }
    }
    // Wait for connection be established or failed
    val device = ProxyDevice(input, output, doClose = {
        // we need to explicitly close the coroutine job, or it can hang for a long time
        // leaking resources.
        runCatching { output.close() }
        runCatching { input.close() }
        closeHandle.complete(true)
//            job.cancel()
    })
    return device
}