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.Protocol
import okhttp3.Request
import okhttp3.Response
import okhttp3.WebSocketListener
import okio.Buffer
import okio.ByteString
import org.java_websocket.WebSocket
import java.io.IOException
import java.net.InetAddress
import java.net.Socket
import java.security.KeyStore
import java.util.concurrent.CountDownLatch
import java.util.concurrent.TimeUnit
import java.util.concurrent.atomic.AtomicBoolean
import java.util.concurrent.atomic.AtomicInteger
import java.util.concurrent.atomic.AtomicLong
import java.util.concurrent.atomic.AtomicReference
import javax.net.SocketFactory
import javax.net.ssl.SSLContext
import javax.net.ssl.TrustManagerFactory
import javax.net.ssl.X509TrustManager


internal const val MAX_STAGED_FRAMES = 1024

/**
 * 생성하는 모든 plain socket에 `tcpNoDelay = true`를 적용하는 [SocketFactory].
 * OkHttp의 SSL socket factory는 이 plain socket 위에 layering되므로 ws/wss 모두 no-delay가 된다.
 */
internal object NoDelaySocketFactory : SocketFactory() {
    private val delegate: SocketFactory = getDefault()

    private fun configure(socket: Socket): Socket = socket.apply { tcpNoDelay = true }

    override fun createSocket(): Socket = configure(delegate.createSocket())

    override fun createSocket(host: String?, port: Int): Socket =
        configure(delegate.createSocket(host, port))

    override fun createSocket(host: String?, port: Int, localHost: InetAddress?, localPort: Int): Socket =
        configure(delegate.createSocket(host, port, localHost, localPort))

    override fun createSocket(host: InetAddress?, port: Int): Socket =
        configure(delegate.createSocket(host, port))

    override fun createSocket(host: InetAddress?, port: Int, localAddress: InetAddress?, localPort: Int): Socket =
        configure(delegate.createSocket(host, port, localAddress, localPort))
}

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(
        object : Transport {
            override suspend fun writePacket(base: PacketBase) = this@WsClient.writePacket(base)

            override fun close() {
                // BaseClient owns the physical websocket close after Communicator drains state.
            }
        },
        parserMap,
        scope,
    )
    private val receiveQueue = Channel<ByteArray>(capacity = maxStagedFrames)

    internal val metrics = WsMetrics()
    private val stagedCount = AtomicInteger(0)

    internal fun getMetricsSnapshot(): WsMetricsSnapshot = metrics.snapshot()

    /** Returns diagnostic counters as a [LongArray]:
     *  [0]=maxStaged, [1]=droppedCount, [2]=totalCopyNanos, [3]=receivedCount */
    internal fun getMetricsDiagnostics(): LongArray {
        val s = metrics.snapshot()
        return longArrayOf(s.maxStaged, s.droppedCount, s.totalCopyNanos, s.receivedCount)
    }

    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, copyNanos: Long = 0L) {
        if (!communicator.isAlive()) return
        metrics.recordReceived(copyNanos)
        val result = receiveQueue.trySend(bytes)
        if (result.isSuccess) {
            val count = stagedCount.incrementAndGet()
            metrics.updateMaxStaged(count.toLong())
            metrics.recordEnqueue(true)
        } else {
            metrics.recordEnqueue(false)
            onDisconnected()
        }
    }

    private fun receiveLoop() {
        scope.launch {
            for (bytes in receiveQueue) {
                stagedCount.decrementAndGet()
                // 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?>()
    private val closed = AtomicBoolean(false)
    private val closedLatch = CountDownLatch(1)

    fun setSocket(ws: okhttp3.WebSocket) {
        if (closed.get()) {
            runCatching { ws.cancel() }
            return
        }
        socketRef.set(ws)
        if (closed.get() && socketRef.compareAndSet(ws, null)) {
            runCatching { ws.cancel() }
        }
    }

    override fun send(bytes: ByteArray): Boolean =
        socketRef.get()?.send(Buffer().write(bytes).readByteString()) ?: false

    override fun close() {
        closeSocket(cancel = false)
    }

    fun cancel() {
        closeSocket(cancel = true)
    }

    fun markClosed() {
        closedLatch.countDown()
    }

    private fun closeSocket(cancel: Boolean) {
        if (!closed.compareAndSet(false, true)) {
            if (!cancel) {
                awaitClosed()
            }
            return
        }
        val socket = socketRef.getAndSet(null)
        if (socket != null) {
            if (cancel) {
                runCatching { socket.cancel() }
                markClosed()
            } else {
                val enqueued = runCatching { socket.close(1000, null) }.getOrDefault(false)
                if (!enqueued) {
                    runCatching { socket.cancel() }
                    markClosed()
                } else if (!awaitClosed()) {
                    runCatching { socket.cancel() }
                    markClosed()
                }
            }
        } else {
            markClosed()
        }
        client.connectionPool.evictAll()
        client.dispatcher.executorService.shutdown()
    }

    private fun awaitClosed(): Boolean =
        runCatching { closedLatch.await(5, TimeUnit.SECONDS) }
            .getOrElse {
                Thread.currentThread().interrupt()
                false
            }
}

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)
        .protocols(listOf(Protocol.HTTP_1_1))
        .readTimeout(30, java.util.concurrent.TimeUnit.SECONDS)
        .writeTimeout(30, java.util.concurrent.TimeUnit.SECONDS)
        // 서버(WsServer)는 이미 tcpNoDelay를 켠다. 클라이언트 plain socket도 no-delay로 만들어
        // Nagle's algorithm + delayed-ACK 조합에서 생기는 ~40ms 고정 지연을 막는다.
        // SSL socket factory는 이 plain socket 위에 layering되므로 ws/wss 모두에 적용된다.
        .socketFactory(NoDelaySocketFactory)
    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) {
            val start = System.nanoTime()
            val byteArray = bytes.toByteArray()
            val copyNanos = System.nanoTime() - start
            client?.receiveBytes(byteArray, copyNanos)
        }

        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) {
            connection.markClosed()
            client?.handleFailure()
        }

        override fun onFailure(webSocket: okhttp3.WebSocket, t: Throwable, response: Response?) {
            failed.set(t)
            opened.countDown()
            connection.markClosed()
            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)) {
        connection.cancel()
        client.close()
        throw IOException("websocket open timed out")
    }
    failed.get()?.let {
        connection.cancel()
        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) }
}

internal data class WsMetricsSnapshot(
    val receivedCount: Long,
    val enqueuedCount: Long,
    val droppedCount: Long,
    val maxStaged: Long,
    val totalCopyNanos: Long
)

internal class WsMetrics {
    private val received = AtomicLong(0)
    private val enqueued = AtomicLong(0)
    private val dropped = AtomicLong(0)
    private val maxStagedVal = AtomicLong(0)
    private val copyNanos = AtomicLong(0)

    fun recordReceived(copyNanosSpent: Long) {
        received.incrementAndGet()
        copyNanos.addAndGet(copyNanosSpent)
    }

    fun recordEnqueue(success: Boolean) {
        if (success) {
            enqueued.incrementAndGet()
        } else {
            dropped.incrementAndGet()
        }
    }

    fun updateMaxStaged(size: Long) {
        while (true) {
            val current = maxStagedVal.get()
            if (size <= current) break
            if (maxStagedVal.compareAndSet(current, size)) break
        }
    }

    fun snapshot(): WsMetricsSnapshot {
        return WsMetricsSnapshot(
            receivedCount = received.get(),
            enqueuedCount = enqueued.get(),
            droppedCount = dropped.get(),
            maxStaged = maxStagedVal.get(),
            totalCopyNanos = copyNanos.get()
        )
    }
}
