package moe.lava.neon.api.gateway

import co.touchlab.kermit.Logger
import io.ktor.client.HttpClient
import io.ktor.client.plugins.cookies.HttpCookies
import io.ktor.client.plugins.websocket.DefaultClientWebSocketSession
import io.ktor.client.plugins.websocket.WebSockets
import io.ktor.client.plugins.websocket.webSocketSession
import io.ktor.client.request.parameter
import io.ktor.http.userAgent
import io.ktor.websocket.Frame
import io.ktor.websocket.readText
import io.ktor.websocket.send
import kotlinx.coroutines.CoroutineScope
import kotlinx.coroutines.Dispatchers
import kotlinx.coroutines.ExperimentalCoroutinesApi
import kotlinx.coroutines.cancel
import kotlinx.coroutines.delay
import kotlinx.coroutines.flow.consumeAsFlow
import kotlinx.coroutines.flow.launchIn
import kotlinx.coroutines.flow.onCompletion
import kotlinx.coroutines.flow.onEach
import kotlinx.coroutines.isActive
import kotlinx.coroutines.launch
import moe.lava.neon.api.ApiConstants
import moe.lava.neon.api.ApiConstants.json
import kotlin.random.Random
import kotlin.time.Duration.Companion.milliseconds

private val logger = Logger.withTag("neon.core.api.gateway/session")

internal class GatewaySession private constructor(
    private var ws: DefaultClientWebSocketSession,
    private val token: String,
    private val scope: CoroutineScope,
    private var resumeProps: ResumeProperties?,
    private val onDispatch: (Event.Dispatch) -> Unit,
    private val onDestroy: (GatewayCloseReason, ResumeProperties?) -> Unit,
    private val onSuccess: () -> Unit,
) {
    private var lastSeq: Int? = resumeProps?.lastSequence
    private var missedHeartbeats = 0
    private var closeReason: GatewayCloseReason? = null

    companion object {
        suspend fun start(
            token: String,
            client: HttpClient = HttpClient {
                install(HttpCookies)
                install(WebSockets)
            },
            scope: CoroutineScope = CoroutineScope(Dispatchers.IO),
            resumeProps: ResumeProperties? = null,
            onDispatch: (Event.Dispatch) -> Unit,
            onDestroy: (GatewayCloseReason, ResumeProperties?) -> Unit,
            onSuccess: () -> Unit,
        ): GatewaySession {
            val ws = client.webSocketSession(
                resumeProps?.resumeGatewayUrl ?: "wss://gateway.discord.gg/"
            ) {
                userAgent(ApiConstants.gatewayUserAgent)
                url {
                    parameter("encoding", "json")
                    parameter("v", "9")
//                    parameter("compress", "zstd-stream")
                }
            }

            return GatewaySession(ws, token, scope, resumeProps, onDispatch, onDestroy, onSuccess)
        }
    }

    init {
        ws.incoming.consumeAsFlow()
            .onCompletion { onClose(it) }
            .onEach { frame ->
                if (frame !is Frame.Text)
//                if (frame !is Frame.Text && frame !is Frame.Binary)
                    return@onEach

                logger.d { "Received payload ${frame.readText()}" }

                val raw = json.decodeFromString<Payload.Unknown>(frame.readText())
                val seq = this.lastSeq ?: 0
                if (seq + 1 == raw.s) {
                    this.lastSeq = raw.s
                } else if (raw.s != null) {
                    close(GatewayCloseReason.SkippedSequence(raw.s))
                    return@onEach
                }

                when (val payload = raw.asIncoming()) {
                    is Payload.Incoming<*> -> scope.launch { handlePayload(payload) }
                    is Payload.Unknown -> scope.launch { handleUnknownPayload(payload) }
                }
            }
            .launchIn(scope)
    }

    private suspend fun handlePayload(payload: Payload.Incoming<*>) {
        logger.d { payload.toString() }
        val event = payload.d
        when (event) {
            is Event.Heartbeat -> handleHeartbeat()
            is Event.Reconnect -> close(GatewayCloseReason.ServerReconnect)
            is Event.InvalidSession -> close(GatewayCloseReason.InvalidSession(event.resumable))
            is Event.Hello -> handleHello(event)
            is Event.HeartbeatAck -> { missedHeartbeats -= 1 }

            is Event.Ready -> {
                resumeProps = ResumeProperties(
                    sessionId = event.sessionId,
                    resumeGatewayUrl = event.resumeGatewayUrl,
                    lastSequence = 0,
                )
                onSuccess()
            }
            is Event.Resumed -> onSuccess()
        }
        if (event is Event.Dispatch) {
            onDispatch(event)
        }
    }

    private suspend fun handleUnknownPayload(payload: Payload.Unknown) {
        logger.w { "Unknown payload $payload" }
    }

    private suspend fun handleHeartbeat() {
        logger.w { "Received heartbeat from server, possible connection issue" }
        Event.QoSHeartbeat(lastSeq).pack().send()
        missedHeartbeats += 1
    }

    private suspend fun handleHello(e: Event.Hello) {
        val resumeProps = resumeProps
        if (resumeProps != null) {
            Event.Resume(
                token = token,
                sessionId = resumeProps.sessionId,
                seq = resumeProps.lastSequence
            ).pack().send()
        } else {
            Event.Identify(token = token).pack().send()
        }

        val interval = e.heartbeatInterval.milliseconds
        scope.launch {
            delay(interval * Random.nextDouble())
            while (true) {
                if (missedHeartbeats >= 1) {
                    close(GatewayCloseReason.MissedHeartbeat)
                    break
                }
                Event.QoSHeartbeat(lastSeq).pack().send()
                missedHeartbeats += 1
                delay(interval)
            }
        }
    }

    fun close(reason: GatewayCloseReason.ClientInitiated?) {
        val msg = when (reason) {
            is GatewayCloseReason.MissedHeartbeat ->
                "heartbeat missed"
            is GatewayCloseReason.SkippedSequence ->
                "payloads skipped one sequence (expected: $lastSeq, actual: ${reason.next})"
            is GatewayCloseReason.InvalidSession ->
                "invalid session (resumable: $reason)"
            is GatewayCloseReason.ClientPaused ->
                "client requested pause"
            is GatewayCloseReason.ServerReconnect ->
                "server requested reconnect"
            null ->
                "no reason"

        }
        closeReason = reason

        logger.e { "Client-initiated close, cause: $msg" }
        ws.cancel()
    }

    @OptIn(ExperimentalCoroutinesApi::class)
    private fun onClose(error: Throwable? = null) {
        logger.d(error) { "Websocket connection closed, cleaning up..." }
        if (scope.isActive) scope.cancel()
        if (resumeProps == null) {
            logger.w { "No resume props stored" }
        }
        onDestroy(
            closeReason
                ?: runCatching { ws.closeReason.getCompleted() }
                    .getOrNull()
                    ?.let { GatewayCloseReason.ServerClosed(it) }
                ?: GatewayCloseReason.Unknown,
            resumeProps?.copy(lastSequence = lastSeq ?: 0)
        )
    }

    private suspend inline fun <reified T : Event.Outgoing> Payload.Outgoing<T>.send() {
        logger.d { "Sending payload $this" }
        logger.d { "Raw: ${json.encodeToString(this)}" }
        ws.send(json.encodeToString(this))
    }
}
