package com.tokilabs.proto_socket

import com.google.protobuf.MessageLite
import kotlinx.coroutines.runBlocking
import org.java_websocket.WebSocket
import org.java_websocket.handshake.ClientHandshake
import org.java_websocket.server.DefaultSSLWebSocketServerFactory
import org.java_websocket.server.WebSocketServer
import java.io.IOException
import java.net.InetSocketAddress
import java.nio.ByteBuffer
import java.nio.channels.SelectionKey
import java.nio.channels.SocketChannel
import java.util.concurrent.ConcurrentHashMap
import java.util.concurrent.CopyOnWriteArrayList
import java.util.concurrent.CountDownLatch
import java.util.concurrent.TimeUnit
import java.util.concurrent.atomic.AtomicBoolean
import javax.net.ssl.SSLContext

class WsServer(
    host: String,
    port: Int,
    private val path: String = "/",
    private val sslContext: SSLContext? = null,
    private val newClient: (WebSocket) -> WsClient,
) : WebSocketServer(InetSocketAddress(host, port)) {
    private val clientByConn = ConcurrentHashMap<WebSocket, WsClient>()
    private val clients = CopyOnWriteArrayList<WsClient>()
    private val startedFlag = AtomicBoolean(false)
    private val startLatch = CountDownLatch(1)
    @Volatile
    private var startError: Exception? = null

    var onClientConnected: (WsClient) -> Unit = {}

    init {
        setReuseAddr(true)
        if (sslContext != null) {
            setWebSocketFactory(DefaultSSLWebSocketServerFactory(sslContext))
        }
    }

    fun started(): Boolean = startedFlag.get()

    fun clients(): List<WsClient> = clients.toList()

    fun port(): Int = getPort()

    override fun onConnect(key: SelectionKey): Boolean {
        (key.channel() as? SocketChannel)?.socket()?.tcpNoDelay = true
        return super.onConnect(key)
    }

    override fun start() {
        if (!startedFlag.compareAndSet(false, true)) return
        try {
            super.start()
        } catch (error: RuntimeException) {
            startedFlag.set(false)
            throw error
        }
        if (!startLatch.await(5, TimeUnit.SECONDS)) {
            startedFlag.set(false)
            throw IOException("websocket server start timed out")
        }
        val error = startError
        if (error != null) {
            startedFlag.set(false)
            throw IOException("websocket server failed to start", error)
        }
    }

    override fun stop() {
        if (!startedFlag.compareAndSet(true, false)) return
        val snapshot = clients.toList()
        snapshot.forEach { it.close() }
        runCatching { super.stop(1000) }
        clients.clear()
        clientByConn.clear()
    }

    suspend fun broadcast(message: MessageLite) {
        clients().forEach { client ->
            client.communicator.send(message)
        }
    }

    override fun onOpen(conn: WebSocket, handshake: ClientHandshake) {
        if (handshake.resourceDescriptor != path) {
            conn.close(1008, "unexpected path")
            return
        }
        val client = newClient(conn)
        client.addDisconnectListener {
            clients.remove(it)
            clientByConn.remove(conn)
        }
        clientByConn[conn] = client
        clients.add(client)
        onClientConnected(client)
    }

    override fun onMessage(conn: WebSocket, message: String) {
    }

    override fun onMessage(conn: WebSocket, message: ByteBuffer) {
        val start = System.nanoTime()
        val bytes = ByteArray(message.remaining())
        message.get(bytes)
        val copyNanos = System.nanoTime() - start
        clientByConn[conn]?.receiveBytes(bytes, copyNanos)
    }

    override fun onClose(
        conn: WebSocket,
        code: Int,
        reason: String,
        remote: Boolean,
    ) {
        removeClient(conn)?.handleFailure()
    }

    override fun onError(conn: WebSocket?, ex: Exception) {
        if (conn == null) {
            if (startLatch.count > 0) {
                startError = ex
                startLatch.countDown()
            }
            return
        }
        removeClient(conn)?.handleFailure()
    }

    override fun onStart() {
        startLatch.countDown()
    }

    private fun removeClient(conn: WebSocket): WsClient? {
        val client = clientByConn.remove(conn) ?: return null
        clients.remove(client)
        return client
    }
}

fun WsServer.broadcastBlocking(message: MessageLite) {
    runBlocking { broadcast(message) }
}
