package com.tokilabs.proto_socket

import com.tokilabs.proto_socket.packets.HeartBeat
import kotlinx.coroutines.CompletableDeferred
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 kotlinx.coroutines.sync.Mutex
import kotlinx.coroutines.sync.withLock
import java.util.concurrent.CopyOnWriteArrayList
import java.util.concurrent.atomic.AtomicBoolean

abstract class BaseClient<Self : BaseClient<Self>>(
    intervalSec: Int,
    waitSec: Int,
    private val doClose: () -> Unit,
    internal val scope: CoroutineScope = CoroutineScope(SupervisorJob() + Dispatchers.IO),
) {
    abstract val self: Self
    abstract val communicator: Communicator

    private val heartbeatIntervalMs = intervalSec * 1000L
    private val heartbeatWaitMs = waitSec * 1000L
    private val closedOnce = AtomicBoolean(false)
    private val closedSignal = CompletableDeferred<Unit>()
    private val hbMutex = Mutex()
    private var hbTimer: HeartbeatTimer? = null
    private var waitingHBResponse = false
    private val disconnectListeners = CopyOnWriteArrayList<(Self) -> Unit>()

    fun addDisconnectListener(handler: (Self) -> Unit) {
        disconnectListeners.add(handler)
    }

    fun removeDisconnectListeners() {
        disconnectListeners.clear()
    }

    open fun close() {
        runBlocking {
            closeWithGrace(true)
            closedSignal.await()
        }
    }

    suspend fun closeAsync() {
        closeWithGrace(true)
        closedSignal.await()
    }

    private suspend fun closeWithGrace(graceful: Boolean) {
        if (!closedOnce.compareAndSet(false, true)) return
        try {
            if (graceful) {
                communicator.close()
            } else {
                communicator.shutdown()
            }
            stopHeartbeat()
            runCatching { doClose() }
            notifyDisconnected()
        } finally {
            scope.cancel()
            closedSignal.complete(Unit)
        }
    }

    fun sendHeartBeat() {
        if (!communicator.isAlive() || heartbeatIntervalMs <= 0) return
        scope.launch {
            hbMutex.withLock {
                hbTimer?.stop()
                hbTimer = HeartbeatTimer(scope, heartbeatIntervalMs) {
                    if (!communicator.isAlive()) return@HeartbeatTimer
                    hbMutex.withLock {
                        waitingHBResponse = true
                    }
                    runCatching { communicator.send(HeartBeat.getDefaultInstance()) }
                    hbMutex.withLock {
                        hbTimer?.stop()
                        hbTimer = HeartbeatTimer(scope, heartbeatWaitMs) {
                            if (communicator.isAlive()) onDisconnected()
                        }
                    }
                }
            }
        }
    }

    fun onHeartBeat() {
        scope.launch {
            val wasWaiting = hbMutex.withLock {
                if (waitingHBResponse) {
                    waitingHBResponse = false
                    true
                } else {
                    false
                }
            }
            if (!wasWaiting) {
                runCatching { communicator.send(HeartBeat.getDefaultInstance()) }
            }
        }
    }

    fun stopHeartbeat() {
        scope.launch {
            hbMutex.withLock {
                hbTimer?.stop()
                hbTimer = null
            }
        }
    }

    fun onDisconnected() {
        scope.launch {
            closeWithGrace(false)
        }
    }

    private fun notifyDisconnected() {
        val listeners = disconnectListeners.toList()
        disconnectListeners.clear()
        listeners.forEach { it(self) }
    }
}
