package com.tokilabs.proto_socket

import com.google.protobuf.Message
import com.google.protobuf.MessageLite
import com.tokilabs.proto_socket.packets.HeartBeat
import com.tokilabs.proto_socket.packets.PacketBase
import kotlinx.coroutines.CompletableDeferred
import kotlinx.coroutines.CoroutineScope
import kotlinx.coroutines.channels.Channel
import kotlinx.coroutines.channels.ClosedSendChannelException
import kotlinx.coroutines.joinAll
import kotlinx.coroutines.launch
import kotlinx.coroutines.runBlocking
import kotlinx.coroutines.selects.select
import kotlinx.coroutines.withTimeout
import java.util.concurrent.ConcurrentHashMap
import java.util.concurrent.atomic.AtomicBoolean
import java.util.concurrent.atomic.AtomicInteger
import java.util.concurrent.atomic.AtomicLong
import java.util.concurrent.locks.ReentrantReadWriteLock
import kotlin.concurrent.read
import kotlin.concurrent.write
import kotlin.reflect.KClass

class NotConnectedException : IllegalStateException("not connected")

interface Transport {
    suspend fun writePacket(base: PacketBase)

    fun close()
}

typealias ParserMap = Map<String, (ByteArray) -> MessageLite>

private typealias RequestHandler = suspend (MessageLite, Int) -> Unit

internal const val MAX_NONCE: Int = Int.MAX_VALUE

private data class PendingRequest(
    val expectedTypeName: String,
    val ch: Channel<MessageLite> = Channel(capacity = 1),
    val errCh: Channel<Throwable> = Channel(capacity = 1),
)

private data class QueuedPacket(
    val base: PacketBase,
    val done: CompletableDeferred<Throwable?> = CompletableDeferred(),
)

private data class InboundItem(
    val typeName: String,
    val data: ByteArray,
    val incomingNonce: Int,
    val responseNonce: Int,
)

fun typeNameOf(m: MessageLite): String =
    (m as Message).descriptorForType.fullName

@PublishedApi
internal fun <T : Message> defaultMessage(kClass: KClass<T>): T {
    val method = kClass.java.getDeclaredMethod("getDefaultInstance")
    @Suppress("UNCHECKED_CAST")
    return method.invoke(null) as T
}

inline fun <reified T : Message> typeNameOf(): String =
    typeNameOf(defaultMessage(T::class))

class Communicator(
    private val transport: Transport,
    parserMap: ParserMap,
    private val scope: CoroutineScope,
) {
    private val rwLock = ReentrantReadWriteLock()
    private val nonce = AtomicInteger(0)
    private val frameSeq = AtomicLong(0)
    private val isAlive = AtomicBoolean(false)
    private val parserMap: MutableMap<String, (ByteArray) -> MessageLite> = ConcurrentHashMap()
    private val handlers = mutableMapOf<String, MutableList<(MessageLite) -> Unit>>()
    private val reqHandlers = mutableMapOf<String, RequestHandler>()
    private val pendingRequests = mutableMapOf<Int, PendingRequest>()
    private val writeQueue = Channel<QueuedPacket>(capacity = 64)
    private val inboundQueue = Channel<InboundItem>(capacity = 64)
    private val closed = CompletableDeferred<Unit>()
    private val shutdownStarted = AtomicBoolean(false)
    private val isClosing = AtomicBoolean(false)
    @Volatile
    private var receiveLoopJob: kotlinx.coroutines.Job? = null
    private var writeErrorHandler: ((Throwable) -> Unit)? = null
    private var frameErrorHandler: ((Throwable) -> Unit)? = null
    private var gateway: InboundGateway? = null
    private val canonicalNameMap = ConcurrentHashMap<String, String>()

    init {
        initialize(parserMap)
    }

    private fun getShortTypeName(typeName: String): String {
        val lastDot = typeName.lastIndexOf('.')
        return if (lastDot < 0 || lastDot == typeName.length - 1) {
            typeName
        } else {
            typeName.substring(lastDot + 1)
        }
    }

    private fun canonicalize(typeName: String): String {
        return canonicalNameMap[typeName]
            ?: canonicalNameMap[getShortTypeName(typeName)]
            ?: typeName
    }

    private fun initialize(parserMap: ParserMap) {
        val shortToFull = mutableMapOf<String, String>()
        val tempParsers = parserMap.toMutableMap()
        tempParsers[typeNameOf<HeartBeat>()] = { HeartBeat.parseFrom(it) }

        for (fullName in tempParsers.keys) {
            val shortName = getShortTypeName(fullName)
            val existing = shortToFull[shortName]
            if (existing != null && existing != fullName) {
                throw IllegalArgumentException("Duplicate alias mapping: $shortName maps to both $existing and $fullName")
            }
            shortToFull[shortName] = fullName
        }

        for (fullName in tempParsers.keys) {
            canonicalNameMap[fullName] = fullName
            val shortName = getShortTypeName(fullName)
            canonicalNameMap[shortName] = fullName
        }

        for ((k, v) in tempParsers) {
            this.parserMap[canonicalize(k)] = v
        }

        isAlive.set(true)
        isClosing.set(false)
        scope.launch { writeLoop() }
        receiveLoopJob = scope.launch { receiveLoop() }
    }

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

    fun setWriteErrorHandler(fn: (Throwable) -> Unit) {
        rwLock.write { writeErrorHandler = fn }
    }

    /**
     * inbound gateway가 raw frame decode에 실패했을 때 호출되는 handler를 등록한다. transport는 이를 통해
     * gateway path에서도 inline path와 동일한 parse-error disconnect semantics를 적용한다(write path의
     * [setWriteErrorHandler]와 대응). gateway가 없으면 [onReceivedFrame]이 decode 오류를 반환값으로
     * 노출하므로 이 handler는 쓰이지 않는다.
     */
    fun setFrameErrorHandler(fn: (Throwable) -> Unit) {
        rwLock.write { frameErrorHandler = fn }
    }

    private fun onFrameError(err: Throwable) {
        val handler = rwLock.read { frameErrorHandler }
        handler?.invoke(err)
    }

    /**
     * receive coordinator 앞단에 default coroutine worker-pool gateway를 attach한다. [onReceivedFrame]으로
     * 제출된 raw frame은 pool이 decode하고 내부 `seq`로 reorder해 coordinator에 전달하므로 FIFO dispatch와
     * stateful 처리 소유권은 coordinator에 유지된다.
     *
     * gateway는 순수 `PacketBase` envelope decode만 수행하고 pending-request map, listener, write queue를
     * 건드리지 않는다. opt-in이며, attach하지 않으면 [onReceivedFrame]은 호출 coroutine에서 inline decode한다.
     */
    fun enableInboundGateway(workers: Int = 0, queueSize: Int = 64) {
        val gw = CoroutineWorkerGateway(
            scope = scope,
            workers = workers,
            queueSize = queueSize,
            sink = { f -> enqueueInbound(f.typeName, f.data, f.incomingNonce, f.responseNonce) },
            onError = { err ->
                // gateway decode 오류를 transport의 frame error handler로 노출해 inline path와 동일한
                // parse-error disconnect semantics를 적용한다. gateway 자체는 transport를 종료하지 않는다.
                onFrameError(err)
            },
        )
        attachInboundGateway(gw)
    }

    /**
     * custom [InboundGateway]를 receive coordinator 앞단에 설치한다. gateway sink는 coordinator 소유권을
     * 보존하기 위해 decoded frame을 [enqueueInbound]로 전달해야 한다. default worker-pool gateway는
     * [enableInboundGateway]를 사용한다.
     */
    fun attachInboundGateway(gateway: InboundGateway) {
        rwLock.write { this.gateway = gateway }
    }

    private fun closeGateway() {
        val gw = rwLock.write {
            val current = gateway
            gateway = null
            current
        }
        gw?.close()
    }

    /**
     * transport read loop가 raw `PacketBase` frame을 ingest한다.
     *
     * gateway가 attach되어 있으면 frame에 내부 `seq`를 부여해 off-coordinator decode + reorder로 제출하고
     * null을 반환한다. decode에 실패한 frame은 `seq` 순서에 맞춰 비동기로 [setFrameErrorHandler]에 등록된
     * frame error handler에 보고되어 transport가 inline path와 동일한 parse-error disconnect semantics를
     * 적용할 수 있게 한다. gateway 자체는 transport를 종료하지 않는다.
     *
     * gateway가 없으면 호출 coroutine에서 inline decode 후 [onReceivedData]로 전달하고, decode 오류를
     * 반환해 transport가 기존 parse-error disconnect semantics를 적용하게 한다.
     */
    suspend fun onReceivedFrame(raw: ByteArray): Throwable? {
        if (!isAlive() || isClosing.get()) return null
        val gw = rwLock.read { gateway }
        if (gw != null) {
            gw.submit(InboundFrame(seq = frameSeq.incrementAndGet(), bytes = raw))
            return null
        }
        return try {
            val base = PacketBase.parseFrom(raw)
            enqueueInbound(base.typeName, base.data.toByteArray(), base.nonce, base.responseNonce)
            null
        } catch (e: Throwable) {
            e
        }
    }

    @PublishedApi
    internal fun nextNonce(): Int {
        var attempts = 0
        while (attempts < MAX_NONCE) {
            val current = nonce.get()
            val next = if (current >= MAX_NONCE) 1 else current + 1
            if (nonce.compareAndSet(current, next)) {
                val isPending = rwLock.read { pendingRequests.containsKey(next) }
                if (!isPending) {
                    return next
                }
                attempts++
            }
        }
        throw IllegalStateException("All nonces are currently pending")
    }

    fun shutdown() {
        if (!shutdownStarted.compareAndSet(false, true)) return
        isAlive.set(false)
        isClosing.set(false)
        closeGateway()
        writeQueue.close()
        inboundQueue.cancel()
        closed.complete(Unit)

        val pending = rwLock.write {
            val snapshot = pendingRequests.values.toList()
            pendingRequests.clear()
            snapshot
        }
        pending.forEach { it.errCh.trySend(NotConnectedException()) }
    }

    suspend fun close() {
        if (!shutdownStarted.compareAndSet(false, true)) return
        isClosing.set(true)
        closeGateway()
        inboundQueue.close()
        receiveLoopJob?.join()
        
        isAlive.set(false)
        isClosing.set(false)
        writeQueue.close()
        closed.complete(Unit)

        val pending = rwLock.write {
            val snapshot = pendingRequests.values.toList()
            pendingRequests.clear()
            snapshot
        }
        pending.forEach { it.errCh.trySend(NotConnectedException()) }
        transport.close()
    }

    suspend fun queuePacket(base: PacketBase) {
        if (!isAlive()) throw NotConnectedException()
        val item = QueuedPacket(base)
        try {
            select<Unit> {
                writeQueue.onSend(item) {}
                closed.onAwait { throw NotConnectedException() }
            }
        } catch (e: ClosedSendChannelException) {
            throw NotConnectedException()
        }

        val error = select {
            item.done.onAwait { it }
            closed.onAwait { NotConnectedException() }
        }
        if (error != null) throw error
    }

    suspend fun send(m: MessageLite) {
        if (!isAlive()) throw NotConnectedException()
        queuePacket(
            PacketBase.newBuilder()
                .setTypeName(typeNameOf(m))
                .setNonce(nextNonce())
                .setData(m.toByteString())
                .build(),
        )
    }

    suspend fun sendRequest(
        req: MessageLite,
        resTypeName: String,
        timeoutMs: Long = 30_000L,
    ): MessageLite {
        if (!isAlive()) throw NotConnectedException()
        val requestNonce = nextNonce()
        val pending = PendingRequest(expectedTypeName = resTypeName)

        rwLock.write { pendingRequests[requestNonce] = pending }

        try {
            queuePacket(
                PacketBase.newBuilder()
                    .setTypeName(typeNameOf(req))
                    .setNonce(requestNonce)
                    .setData(req.toByteString())
                    .build(),
            )
        } catch (e: Throwable) {
            removePending(requestNonce)
            throw e
        }

        return try {
            withTimeout(if (timeoutMs > 0) timeoutMs else 30_000L) {
                select {
                    pending.ch.onReceive { it }
                    pending.errCh.onReceive { throw it }
                    closed.onAwait { throw NotConnectedException() }
                }
            }
        } catch (e: kotlinx.coroutines.TimeoutCancellationException) {
            removePending(requestNonce)
            throw IllegalStateException("request timeout for nonce $requestNonce", e)
        }
    }

    fun addListener(typeName: String, fn: (MessageLite) -> Unit) {
        rwLock.write {
            val canonical = canonicalize(typeName)
            check(!reqHandlers.containsKey(canonical)) {
                "type $typeName is already registered with addRequestListener"
            }
            handlers.getOrPut(canonical) { mutableListOf() }.add(fn)
        }
    }

    fun removeListeners(typeName: String) {
        rwLock.write {
            val canonical = canonicalize(typeName)
            handlers.remove(canonical)
        }
    }

    fun addRequestListener(typeName: String, fn: RequestHandler) {
        rwLock.write {
            val canonical = canonicalize(typeName)
            check(handlers[canonical].isNullOrEmpty()) {
                "type $typeName is already registered with addListener"
            }
            check(!reqHandlers.containsKey(canonical)) {
                "type $typeName is already registered with addRequestListener"
            }
            reqHandlers[canonical] = fn
        }
    }

    fun onReceivedData(
        typeName: String,
        data: ByteArray,
        incomingNonce: Int = 0,
        responseNonce: Int = 0,
    ) {
        if (responseNonce > 0) {
            handleResponse(typeName, data, responseNonce)
            return
        }
        if (!isAlive() || isClosing.get()) return
        val item = InboundItem(typeName, data, incomingNonce, responseNonce)
        try {
            runBlocking {
                inboundQueue.send(item)
            }
        } catch (e: ClosedSendChannelException) {
            // closed
        }
    }

    suspend fun enqueueInbound(
        typeName: String,
        data: ByteArray,
        incomingNonce: Int = 0,
        responseNonce: Int = 0,
    ) {
        if (responseNonce > 0) {
            handleResponse(typeName, data, responseNonce)
            return
        }
        if (!isAlive() || isClosing.get()) return
        try {
            inboundQueue.send(InboundItem(typeName, data, incomingNonce, responseNonce))
        } catch (e: ClosedSendChannelException) {
            // closed
        }
    }

    private suspend fun writeLoop() {
        for (item in writeQueue) {
            val error = runCatching { transport.writePacket(item.base) }.exceptionOrNull()
            item.done.complete(error)
            if (error != null) {
                val handler = rwLock.read { writeErrorHandler }
                handler?.invoke(error)
            }
        }
    }

    private suspend fun receiveLoop() {
        for (item in inboundQueue) {
            dispatchInbound(item)
        }
    }

    private suspend fun dispatchInbound(item: InboundItem) {
        if (item.responseNonce > 0) {
            handleResponse(item.typeName, item.data, item.responseNonce)
            return
        }

        val reqHandler: RequestHandler?
        val listeners: List<(MessageLite) -> Unit>
        val canonical = canonicalize(item.typeName)
        rwLock.read {
            reqHandler = reqHandlers[canonical]
            listeners = handlers[canonical]?.toList().orEmpty()
        }

        if (reqHandler != null) {
            val msg = runCatching { parse(item.typeName, item.data) }.getOrNull() ?: return
            reqHandler(msg, item.incomingNonce)
            return
        }

        if (listeners.isEmpty()) return
        val msg = runCatching { parse(item.typeName, item.data) }.getOrNull() ?: return
        listeners.forEach { it(msg) }
    }

    private fun handleResponse(
        typeName: String,
        data: ByteArray,
        responseNonce: Int,
    ) {
        val pending = removePending(responseNonce) ?: return
        val canonicalExpected = canonicalize(pending.expectedTypeName)
        val canonicalReceived = canonicalize(typeName)
        if (canonicalReceived != canonicalExpected) {
            pending.errCh.trySend(
                IllegalStateException(
                    "response type mismatch for nonce $responseNonce: " +
                        "expected $canonicalExpected, got $typeName",
                ),
            )
            return
        }
        val msg = runCatching { parse(typeName, data) }
        msg.onSuccess { pending.ch.trySend(it) }
            .onFailure { pending.errCh.trySend(it) }
    }

    internal fun parse(typeName: String, data: ByteArray): MessageLite {
        val canonical = canonicalize(typeName)
        val parser = parserMap[canonical]
            ?: throw IllegalStateException("protobuf parser is not registered for type $typeName")
        return parser(data)
    }

    private fun removePending(nonce: Int): PendingRequest? =
        rwLock.write { pendingRequests.remove(nonce) }
}

/**
 * 수신 측이 gateway 제출 전에 부여하는 내부 raw 프레임이다.
 *
 * `seq`는 단조 증가하는 내부 시퀀스로, 병렬 decode 후 입력 순서를 복원하는 용도로만 쓰인다.
 * wire format이나 공개 프로토콜 필드에는 절대 노출하지 않는다.
 */
data class InboundFrame(val seq: Long, val bytes: ByteArray)

/**
 * gateway가 PacketBase envelope을 decode한 결과다. `Communicator`의 inbound dispatch가
 * 기대하는 envelope 필드와, 입력 순서 복원을 위한 원본 `seq`를 함께 싣는다.
 */
data class DecodedFrame(
    val seq: Long,
    val typeName: String,
    val data: ByteArray,
    val incomingNonce: Int,
    val responseNonce: Int,
)

/**
 * 수신 coordinator 앞단에서 순수 raw-bytes 작업(PacketBase envelope decode/경량 전처리)만
 * 병렬 처리하는 worker gateway 계약이다.
 *
 * gateway는 stateful dispatch, `pendingRequests`, listener, write queue를 소유하지 않는다.
 * decode 결과는 입력 `seq` 순서로 sink에 전달되어 coordinator가 결정적으로 dispatch할 수 있게 한다.
 */
interface InboundGateway {
    /** raw 프레임을 gateway에 제출한다. fire-and-forget이며 decode 결과는 이후 sink에 `seq` 순서로 전달된다. */
    suspend fun submit(frame: InboundFrame)

    /** gateway 자원을 해제하고 추가 전달을 멈춘다. */
    fun close()
}

/** reorder 버퍼를 통과하는 per-frame decode 결과다. 실패 프레임도 seq를 진행시켜 reorder window를 막지 않는다. */
internal data class GatewayResult(
    val seq: Long,
    val frame: DecodedFrame?,
    val error: Throwable?,
)

/**
 * out-of-order로 도착할 수 있는 결과에서 입력 순서를 복원한다. seq는 1부터 연속이라고 가정한다.
 * 결과는 앞선 모든 seq가 release될 때까지 버퍼되며, 이후 자신과 연속된 후속 결과까지 순서대로 반환한다.
 * concurrent 사용에 안전하지 않으며 collector coroutine 단독에서만 호출한다.
 */
internal class FrameReorderBuffer {
    private var nextSeq: Long = 1
    private val pending = HashMap<Long, GatewayResult>()

    /** res를 버퍼하고 지금 순서대로 release 가능한 결과들을 seq 순으로 반환한다. */
    fun release(res: GatewayResult): List<GatewayResult> {
        if (res.seq < nextSeq) {
            // 이미 release되었거나 중복/지연 도착이므로 두 번 방출하지 않도록 무시한다.
            return emptyList()
        }
        pending[res.seq] = res
        val ready = mutableListOf<GatewayResult>()
        while (true) {
            val r = pending.remove(nextSeq) ?: break
            ready.add(r)
            nextSeq++
        }
        return ready
    }
}

/**
 * coroutine worker-pool gateway다. bounded jobs `Channel`이 N개 worker에 프레임을 공급해 PacketBase
 * envelope을 병렬 decode하고, 단일 collector coroutine이 결과를 `seq`로 reorder해 입력 순서로 sink에 전달한다.
 * sink를 호출하는 coroutine이 collector 하나뿐이므로 sink 호출은 직렬화되고 순서가 보존되어, stateful dispatch는
 * 수신 coordinator에 남는다.
 *
 * decode에 실패한 프레임은 sink가 아니라 `onError`로 `seq` 순서에 맞춰 보고되어 transport가 기존 parse-error
 * 정책을 적용할 수 있게 한다. gateway 자체는 transport를 종료하지 않는다. `onError`는 null일 수 있다.
 *
 * `workers <= 0`은 사용 가능한 프로세서 수로, `queueSize <= 0`은 64로 기본값을 적용한다. bounded jobs
 * 채널이 backpressure를 제공한다. 채널이 가득 차면 worker가 슬롯을 비울 때까지 `submit`이 suspend한다.
 */
class CoroutineWorkerGateway(
    scope: CoroutineScope,
    workers: Int = 0,
    queueSize: Int = 64,
    private val decode: suspend (ByteArray) -> PacketBase = { PacketBase.parseFrom(it) },
    private val sink: suspend (DecodedFrame) -> Unit,
    private val onError: ((Throwable) -> Unit)? = null,
) : InboundGateway {
    private val jobs: Channel<InboundFrame>
    private val results: Channel<GatewayResult>
    private val closed = AtomicBoolean(false)

    init {
        val workerCount = if (workers <= 0) Runtime.getRuntime().availableProcessors() else workers
        val capacity = if (queueSize <= 0) 64 else queueSize
        jobs = Channel(capacity = capacity)
        results = Channel(capacity = capacity)

        val workerJobs = (1..workerCount).map { scope.launch { workerLoop() } }
        scope.launch { collectorLoop() }
        // 모든 worker가 jobs를 drain하고 종료하면 results를 닫아 collector도 종료시킨다.
        scope.launch {
            workerJobs.joinAll()
            results.close()
        }
    }

    private suspend fun workerLoop() {
        for (frame in jobs) {
            val result = try {
                val base = decode(frame.bytes)
                GatewayResult(
                    seq = frame.seq,
                    frame = DecodedFrame(
                        seq = frame.seq,
                        typeName = base.typeName,
                        data = base.data.toByteArray(),
                        incomingNonce = base.nonce,
                        responseNonce = base.responseNonce,
                    ),
                    error = null,
                )
            } catch (e: Throwable) {
                GatewayResult(seq = frame.seq, frame = null, error = e)
            }
            try {
                results.send(result)
            } catch (e: ClosedSendChannelException) {
                return
            }
        }
    }

    private suspend fun collectorLoop() {
        val reorder = FrameReorderBuffer()
        for (result in results) {
            for (ready in reorder.release(result)) {
                val error = ready.error
                if (error != null) {
                    // decode 오류를 seq 순서로 보고하고 계속 진행해 한 프레임의 실패가 이후 프레임을 막지 않게 한다.
                    // transport 종료 여부는 호출자의 onError가 결정한다.
                    onError?.invoke(error)
                    continue
                }
                ready.frame?.let { sink(it) }
            }
        }
    }

    /** raw 프레임을 worker pool에 넘긴다. bounded jobs 채널이 가득 차면 suspend(backpressure)하고, 닫혀 있으면 즉시 drop한다. */
    override suspend fun submit(frame: InboundFrame) {
        if (closed.get()) return
        try {
            jobs.send(frame)
        } catch (e: ClosedSendChannelException) {
            // gateway가 닫혔으므로 drop한다.
        }
    }

    /** jobs 채널을 닫아 worker가 남은 프레임을 drain한 뒤 종료하게 한다. 멱등하며 in-flight sink 호출을 기다리지 않는다. */
    override fun close() {
        if (!closed.compareAndSet(false, true)) return
        jobs.close()
    }
}

inline fun <reified T : Message> addListenerTyped(
    communicator: Communicator,
    noinline fn: (T) -> Unit,
) {
    val typeName = typeNameOf<T>()
    communicator.addListener(typeName) { message ->
        val typed = message as? T
            ?: throw IllegalStateException("received ${message::class} for listener $typeName")
        fn(typed)
    }
}

inline fun <reified Req : Message, reified Res : Message> addRequestListenerTyped(
    communicator: Communicator,
    noinline fn: suspend (Req) -> Res,
) {
    val reqTypeName = typeNameOf<Req>()
    communicator.addRequestListener(reqTypeName) { message, requestNonce ->
        val req = message as? Req
        if (req != null) {
            val res = fn(req)
            if (communicator.isAlive()) {
                communicator.queuePacket(
                    PacketBase.newBuilder()
                        .setTypeName(typeNameOf(res))
                        .setNonce(communicator.nextNonce())
                        .setResponseNonce(requestNonce)
                        .setData(res.toByteString())
                        .build(),
                )
            }
        }
    }
}

suspend inline fun <reified Req : Message, reified Res : Message> sendRequestTyped(
    communicator: Communicator,
    req: Req,
    timeoutMs: Long = 30_000L,
): Res {
    val msg = communicator.sendRequest(req, typeNameOf<Res>(), timeoutMs)
    return msg as? Res
        ?: throw IllegalStateException("received ${msg::class}, expected ${Res::class}")
}
