package com.tokilabs.proto_socket

import com.tokilabs.proto_socket.packets.HeartBeat
import com.tokilabs.proto_socket.packets.PacketBase
import kotlinx.coroutines.CoroutineDispatcher
import kotlinx.coroutines.Dispatchers
import kotlinx.coroutines.asCoroutineDispatcher
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.net.SocketTimeoutException
import java.nio.ByteBuffer
import java.util.concurrent.Executors
import java.util.concurrent.atomic.AtomicInteger
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(
        object : Transport {
            override suspend fun writePacket(base: PacketBase) = this@TcpClient.writePacket(base)

            override fun close() {
                // BaseClient owns the physical socket close after Communicator drains state.
            }
        },
        parserMap,
        scope,
    )
    private val writeMutex = Mutex()

    init {
        socket.tcpNoDelay = true
        socket.soTimeout = READ_POLL_TIMEOUT_MS
        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() {
        // 연결당 read는 blocking이라 Dispatchers.IO에서 돌리면 parallel clients>=128에서 read 스레드가
        // IO lane을 점유해 write/heartbeat/accept를 starvation시킨다. read 전용 dispatcher로 분리해
        // write 경로와 스레드 풀을 공유하지 않게 하고, soTimeout 기반 poll로 close 시 즉시 루프를 빠져나간다.
        scope.launch(readDispatcher) {
            val input = DataInputStream(socket.getInputStream())
            val header = ByteArray(4)
            while (communicator.isAlive()) {
                try {
                    // header/payload 모두 offset-preserving read로 읽는다. soTimeout poll은 frame 경계와
                    // 무관하게 발생할 수 있으므로, readInt/readFully를 직접 쓰면 frame 일부 소비 후 timeout이
                    // 나면 다음 read가 frame 중간 bytes를 length로 오해해 desync된다. readExact는 같은 frame
                    // read 안에서만 timeout을 poll로 흡수해 offset을 보존한다.
                    if (!readExact(input, header)) {
                        // close(isAlive=false)면 정상 종료, EOF(read<0)면 disconnect.
                        if (communicator.isAlive()) onDisconnected()
                        return@launch
                    }
                    val length = ByteBuffer.wrap(header).int
                    if (length == 0) continue
                    if (length < 0 || length > MAX_PACKET_SIZE) {
                        onDisconnected()
                        return@launch
                    }
                    val bytes = ByteArray(length)
                    if (!readExact(input, bytes)) {
                        if (communicator.isAlive()) onDisconnected()
                        return@launch
                    }
                    // 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
                }
            }
        }
    }

    // target을 가득 채울 때까지 offset을 보존하며 읽는다. soTimeout poll(SocketTimeoutException)은 같은
    // read 안에서 흡수해 frame 경계를 잃지 않는다. close되면(isAlive=false) 즉시 false를 반환해 read 루프가
    // socket close에만 의존하지 않고 빠져나간다. EOF(read<0)도 false를 반환해 disconnect로 이어진다.
    // 반환: target.size만큼 다 읽으면 true, close/EOF면 false.
    private fun readExact(input: DataInputStream, target: ByteArray): Boolean {
        var offset = 0
        while (offset < target.size) {
            if (!communicator.isAlive()) return false
            try {
                val read = input.read(target, offset, target.size - offset)
                if (read < 0) return false
                offset += read
            } catch (_: SocketTimeoutException) {
                continue
            }
        }
        return true
    }

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

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

        // soTimeout poll 주기. blocking readInt가 이 주기마다 SocketTimeoutException으로 깨어나
        // isAlive()를 재확인하므로, close 후 read 루프가 socket close에만 의존하지 않고 즉시 종료된다.
        private const val READ_POLL_TIMEOUT_MS: Int = 1_000

        // 모든 TcpClient read 루프가 공유하는 read 전용 dispatcher. cached thread pool이라 연결 수에
        // 맞춰 read 스레드가 늘고, Dispatchers.IO(write/heartbeat/accept)와 풀을 분리해 starvation을
        // 막는다. daemon 스레드라 JVM 종료를 막지 않으며, soTimeout poll로 idle 스레드는 곧 회수된다.
        private val readThreadCounter = AtomicInteger(0)
        private val readDispatcher: CoroutineDispatcher =
            Executors.newCachedThreadPool { runnable ->
                Thread(runnable, "proto-socket-tcp-read-${readThreadCounter.incrementAndGet()}").apply {
                    isDaemon = true
                }
            }.asCoroutineDispatcher()

        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) }
}
