package com.tokilabs.proto_socket

import com.tokilabs.proto_socket.packets.HeartBeat
import com.tokilabs.proto_socket.packets.PacketBase
import kotlinx.coroutines.channels.Channel
import kotlinx.coroutines.launch
import kotlinx.coroutines.runBlocking
import okhttp3.OkHttpClient
import okhttp3.Request
import okhttp3.Response
import okhttp3.WebSocketListener
import okio.ByteString
import org.java_websocket.WebSocket
import java.io.IOException
import java.security.KeyStore
import java.util.concurrent.CountDownLatch
import java.util.concurrent.TimeUnit
import java.util.concurrent.atomic.AtomicReference
import javax.net.ssl.SSLContext
import javax.net.ssl.TrustManagerFactory
import javax.net.ssl.X509TrustManager

internal const val MAX_STAGED_FRAMES = 1024

class WsClient internal constructor(
    private val connection: WsConnection,
    intervalSec: Int,
    waitSec: Int,
    parserMap: ParserMap,
    private val maxStagedFrames: Int = MAX_STAGED_FRAMES,
) : BaseClient<WsClient>(
    intervalSec = intervalSec,
    waitSec = waitSec,
    doClose = { connection.close() },
),
    Transport {
    override val self: WsClient
        get() = this

    override val communicator: Communicator = Communicator(this, parserMap, scope)
    private val receiveQueue = Channel<ByteArray>(capacity = maxStagedFrames)

    init {
        communicator.setWriteErrorHandler { onDisconnected() }
        communicator.setFrameErrorHandler { onDisconnected() }
        communicator.addListener(typeNameOf<HeartBeat>()) { onHeartBeat() }
        receiveLoop()
        sendHeartBeat()
    }

    override suspend fun writePacket(base: PacketBase) {
        val bytes = base.toByteArray()
        if (!connection.send(bytes)) {
            throw IOException("websocket send failed")
        }
    }

    internal fun receiveBytes(bytes: ByteArray) {
        if (!communicator.isAlive()) return
        val result = receiveQueue.trySend(bytes)
        if (result.isFailure) {
            onDisconnected()
        }
    }

    private fun receiveLoop() {
        scope.launch {
            for (bytes in receiveQueue) {
                // 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()
            }
        }
    }

    internal fun handleFailure() {
        onDisconnected()
    }

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

    companion object {
        fun forServer(
            conn: WebSocket,
            intervalSec: Int,
            waitSec: Int,
            parserMap: ParserMap,
        ): WsClient = WsClient(JavaWebSocketConnection(conn), intervalSec, waitSec, parserMap)
    }
}

internal interface WsConnection {
    fun send(bytes: ByteArray): Boolean

    fun close()
}

private class JavaWebSocketConnection(
    private val conn: WebSocket,
) : WsConnection {
    override fun send(bytes: ByteArray): Boolean =
        runCatching {
            conn.send(bytes)
            true
        }.getOrElse { false }

    override fun close() {
        runCatching { conn.close(1000) }
    }
}

private class OkHttpWsConnection(
    private val client: OkHttpClient,
) : WsConnection {
    private val socketRef = AtomicReference<okhttp3.WebSocket?>()

    fun setSocket(ws: okhttp3.WebSocket) {
        socketRef.set(ws)
    }

    override fun send(bytes: ByteArray): Boolean =
        socketRef.get()?.send(ByteString.of(*bytes)) ?: false

    override fun close() {
        runCatching { socketRef.get()?.close(1000, null) }
        client.dispatcher.executorService.shutdown()
    }
}

fun dialWs(
    host: String,
    port: Int,
    path: String,
    intervalSec: Int = 30,
    waitSec: Int = 10,
    parserMap: ParserMap,
): WsClient = dialWsUrl("ws://$host:$port$path", null, intervalSec, waitSec, parserMap)

fun dialWss(
    host: String,
    port: Int,
    path: String,
    sslContext: SSLContext?,
    intervalSec: Int = 30,
    waitSec: Int = 10,
    parserMap: ParserMap,
): WsClient = dialWsUrl("wss://$host:$port$path", sslContext, intervalSec, waitSec, parserMap)

private fun dialWsUrl(
    url: String,
    sslContext: SSLContext?,
    intervalSec: Int,
    waitSec: Int,
    parserMap: ParserMap,
): WsClient {
    val builder = OkHttpClient.Builder()
        .connectTimeout(30, java.util.concurrent.TimeUnit.SECONDS)
        .readTimeout(30, java.util.concurrent.TimeUnit.SECONDS)
        .writeTimeout(30, java.util.concurrent.TimeUnit.SECONDS)
    if (sslContext != null) {
        val trustManager = object : javax.net.ssl.X509TrustManager {
            override fun checkClientTrusted(chain: Array<out java.security.cert.X509Certificate>?, authType: String?) {}
            override fun checkServerTrusted(chain: Array<out java.security.cert.X509Certificate>?, authType: String?) {}
            override fun getAcceptedIssuers(): Array<java.security.cert.X509Certificate> = arrayOf()
        }
        builder.sslSocketFactory(sslContext.socketFactory, trustManager)
    }
    val okHttpClient = builder.build()
    val connection = OkHttpWsConnection(okHttpClient)
    val opened = CountDownLatch(1)
    val failed = AtomicReference<Throwable?>()
    var client: WsClient? = null
    val listener = object : WebSocketListener() {
        override fun onOpen(webSocket: okhttp3.WebSocket, response: Response) {
            connection.setSocket(webSocket)
            opened.countDown()
        }

        override fun onMessage(webSocket: okhttp3.WebSocket, bytes: ByteString) {
            client?.receiveBytes(bytes.toByteArray())
        }

        override fun onClosing(webSocket: okhttp3.WebSocket, code: Int, reason: String) {
            webSocket.close(code, reason)
            client?.handleFailure()
        }

        override fun onClosed(webSocket: okhttp3.WebSocket, code: Int, reason: String) {
            client?.handleFailure()
        }

        override fun onFailure(webSocket: okhttp3.WebSocket, t: Throwable, response: Response?) {
            failed.set(t)
            opened.countDown()
            client?.handleFailure()
        }
    }

    client = WsClient(connection, intervalSec, waitSec, parserMap)
    val ws = okHttpClient.newWebSocket(Request.Builder().url(url).build(), listener)
    connection.setSocket(ws)
    if (!opened.await(20, TimeUnit.SECONDS)) {
        client.close()
        throw IOException("websocket open timed out")
    }
    failed.get()?.let {
        client.close()
        throw IOException("websocket open failed", it)
    }
    return client
}

private fun defaultTrustManager(): X509TrustManager {
    val factory = TrustManagerFactory.getInstance(TrustManagerFactory.getDefaultAlgorithm())
    factory.init(null as KeyStore?)
    return factory.trustManagers
        .filterIsInstance<X509TrustManager>()
        .singleOrNull()
        ?: throw IllegalStateException("expected exactly one default X509TrustManager")
}

@Suppress("FunctionName")
fun DialWs(
    host: String,
    port: Int,
    path: String,
    intervalSec: Int = 30,
    waitSec: Int = 10,
    parserMap: ParserMap,
): WsClient = dialWs(host, port, path, intervalSec, waitSec, parserMap)

@Suppress("FunctionName")
fun DialWss(
    host: String,
    port: Int,
    path: String,
    sslContext: SSLContext?,
    intervalSec: Int = 30,
    waitSec: Int = 10,
    parserMap: ParserMap,
): WsClient = dialWss(host, port, path, sslContext, intervalSec, waitSec, parserMap)

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