package com.tokilabs.proto_socket

import com.tokilabs.proto_socket.packets.HeartBeat
import com.tokilabs.proto_socket.packets.PacketBase
import kotlinx.coroutines.Dispatchers
import kotlinx.coroutines.launch
import kotlinx.coroutines.runBlocking
import kotlinx.coroutines.sync.Mutex
import kotlinx.coroutines.sync.withLock
import kotlinx.coroutines.withContext
import java.io.DataInputStream
import java.io.IOException
import java.net.InetSocketAddress
import java.net.Socket
import java.nio.ByteBuffer
import javax.net.ssl.SSLContext

class TcpClient private constructor(
    private val socket: Socket,
    intervalSec: Int,
    waitSec: Int,
    parserMap: ParserMap,
) : BaseClient<TcpClient>(
    intervalSec = intervalSec,
    waitSec = waitSec,
    doClose = { socket.close() },
),
    Transport {
    override val self: TcpClient
        get() = this

    override val communicator: Communicator = Communicator(this, parserMap, scope)
    private val writeMutex = Mutex()

    init {
        socket.tcpNoDelay = true
        communicator.setWriteErrorHandler { onDisconnected() }
        communicator.setFrameErrorHandler { onDisconnected() }
        communicator.addListener(typeNameOf<HeartBeat>()) { onHeartBeat() }
        readLoop()
        sendHeartBeat()
    }

    override suspend fun writePacket(base: PacketBase) {
        val bytes = base.toByteArray()
        if (bytes.size > MAX_PACKET_SIZE) {
            throw IOException("packet size ${bytes.size} exceeds max $MAX_PACKET_SIZE")
        }
        val header = ByteBuffer.allocate(4).putInt(bytes.size).array()
        writeMutex.withLock {
            withContext(Dispatchers.IO) {
                val output = socket.getOutputStream()
                output.write(header)
                output.write(bytes)
                output.flush()
            }
        }
    }

    private fun readLoop() {
        scope.launch(Dispatchers.IO) {
            val input = DataInputStream(socket.getInputStream())
            while (communicator.isAlive()) {
                try {
                    val length = input.readInt()
                    if (length == 0) continue
                    if (length < 0 || length > MAX_PACKET_SIZE) {
                        onDisconnected()
                        return@launch
                    }
                    val bytes = ByteArray(length)
                    input.readFully(bytes)
                    // raw frame을 communicator ingress로 넘긴다. gateway가 attach되어 있으면 worker pool이
                    // decode/reorder하고, 없으면 inline decode한다. inline decode 오류는 반환값으로,
                    // gateway decode 오류는 frame error handler로 동일한 disconnect semantics를 적용한다.
                    if (communicator.onReceivedFrame(bytes) != null) {
                        onDisconnected()
                        return@launch
                    }
                    sendHeartBeat()
                } catch (_: Exception) {
                    onDisconnected()
                    return@launch
                }
            }
        }
    }

    suspend fun send(message: com.google.protobuf.MessageLite) {
        communicator.send(message)
    }

    companion object {
        const val MAX_PACKET_SIZE: Int = 64 * 1024 * 1024

        fun fromSocket(
            socket: Socket,
            intervalSec: Int = 30,
            waitSec: Int = 10,
            parserMap: ParserMap,
        ): TcpClient = TcpClient(socket, intervalSec, waitSec, parserMap)
    }
}

fun dialTcp(
    host: String,
    port: Int,
    intervalSec: Int = 30,
    waitSec: Int = 10,
    parserMap: ParserMap,
): TcpClient {
    val socket = Socket()
    socket.connect(InetSocketAddress(host, port))
    return TcpClient.fromSocket(socket, intervalSec, waitSec, parserMap)
}

fun dialTcpTls(
    host: String,
    port: Int,
    sslContext: SSLContext,
    intervalSec: Int = 30,
    waitSec: Int = 10,
    parserMap: ParserMap,
): TcpClient {
    val socket = sslContext.socketFactory.createSocket(host, port)
    return TcpClient.fromSocket(socket, intervalSec, waitSec, parserMap)
}

@Suppress("FunctionName")
fun DialTcp(
    host: String,
    port: Int,
    intervalSec: Int = 30,
    waitSec: Int = 10,
    parserMap: ParserMap,
): TcpClient = dialTcp(host, port, intervalSec, waitSec, parserMap)

@Suppress("FunctionName")
fun DialTcpTls(
    host: String,
    port: Int,
    sslContext: SSLContext,
    intervalSec: Int = 30,
    waitSec: Int = 10,
    parserMap: ParserMap,
): TcpClient = dialTcpTls(host, port, sslContext, intervalSec, waitSec, parserMap)

fun TcpClient.closeBlocking() {
    close()
}

fun TcpClient.sendBlocking(message: com.google.protobuf.MessageLite) {
    runBlocking { send(message) }
}
