package com.tokilabs.proto_socket

import com.google.protobuf.MessageLite
import kotlinx.coroutines.CoroutineScope
import kotlinx.coroutines.Dispatchers
import kotlinx.coroutines.SupervisorJob
import kotlinx.coroutines.cancel
import kotlinx.coroutines.launch
import kotlinx.coroutines.runBlocking
import java.net.InetSocketAddress
import java.net.ServerSocket
import java.net.SocketException
import java.util.concurrent.CopyOnWriteArrayList
import java.util.concurrent.atomic.AtomicBoolean
import javax.net.ssl.SSLContext

class TcpServer(
    private val host: String,
    private val port: Int,
    private val sslContext: SSLContext? = null,
    private val newClient: (java.net.Socket) -> TcpClient,
) {
    private val scope = CoroutineScope(SupervisorJob() + Dispatchers.IO)
    private val clients = CopyOnWriteArrayList<TcpClient>()
    private val started = AtomicBoolean(false)
    private var serverSocket: ServerSocket? = null

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

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

    fun port(): Int = serverSocket?.localPort ?: port

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

    fun start() {
        if (!started.compareAndSet(false, true)) return
        val server: ServerSocket = if (sslContext != null) {
            sslContext.serverSocketFactory.createServerSocket()
        } else {
            ServerSocket()
        }
        server.reuseAddress = true
        server.bind(InetSocketAddress(host, port))
        serverSocket = server
        scope.launch(Dispatchers.IO) {
            acceptLoop(server)
        }
    }

    fun stop() {
        if (!started.compareAndSet(true, false)) return
        runCatching { serverSocket?.close() }
        val snapshot = clients.toList()
        clients.clear()
        snapshot.forEach { it.close() }
        scope.cancel()
    }

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

    private fun acceptLoop(server: ServerSocket) {
        while (started.get()) {
            try {
                val socket = server.accept()
                val client = newClient(socket)
                client.addDisconnectListener { clients.remove(it) }
                clients.add(client)
                onClientConnected(client)
            } catch (_: SocketException) {
                return
            } catch (_: Exception) {
                if (!started.get()) return
            }
        }
    }
}

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