import com.fasterxml.jackson.databind.ObjectMapper
import org.slf4j.LoggerFactory
import org.springframework.web.socket.*
import org.springframework.web.socket.client.WebSocketClient
import org.springframework.web.socket.handler.AbstractWebSocketHandler
import java.util.concurrent.*
import java.util.concurrent.atomic.AtomicBoolean
import java.util.concurrent.atomic.AtomicInteger
/**
 * 包装上游会话及其状态，用于管理授权和心跳
 *
 * @param session WebSocket 上游会话
 * @param authorized 是否已通过授权验证
 * @param downConnected 下游连接是否已建立
 * @param lastHeartbeat 最近心跳时间戳（毫秒）
 */
data class WebSocketProxySession(
    val session: WebSocketSession,
    var authorized: Boolean = false,
    var downConnected: Boolean = false,
    var lastHeartbeat: Long = System.currentTimeMillis()
)
/**
 * 通用 WebSocket 代理抽象类
 *
 * 负责管理上游和下游的连接生命周期、消息转发以及超时清理
 *
 * 使用方式：
 * 1. 实现核心抽象方法：
 *    - registerPath: 定义代理路由
 *    - onUpstreamFirstMessage: 处理上游首条消息并进行授权
 *    - downstreamUri: 获取下游 URI
 *    - transformUpstream: 上游→下游 转换逻辑
 *    - transformDownstream: 下游→上游 转换逻辑
 * 2. 可选覆盖钩子：
 *    - onUpstreamOpen: 上游连接初始化
 *    - onAuthSuccess: 授权成功回调
 *    - onUpstreamFirstMessageIsNull: 授权失败处理
 *    - onSessionClosed: 会话关闭后处理
 *
 * @param objectMapper 用于 JSON 序列化/反序列化
 * @param client WebSocket 客户端，用于建立下游连接
 * @author ThatCoder
 */
abstract class IWebSocketProxier(
    val objectMapper: ObjectMapper,
    private val client: WebSocketClient
) : AbstractWebSocketHandler() {
    /** 代理接入路径 */
    abstract val registerPath: String
    /** 会话超时时间，默认 10 分钟 */
    open val sessionTimeoutMillis: Long = 10 * 60 * 1000
    private val logger = LoggerFactory.getLogger(this::class.java)
    private val sessions = ConcurrentHashMap<String, WebSocketProxySession>()
    private val downstreamContexts = ConcurrentHashMap<String, DownstreamContext>()
    private val scheduler = Executors.newSingleThreadScheduledExecutor(
        NamedThreadFactory("proxy-session-timeout-")
    )
    init {
        // 定期清理超时会话
        scheduler.scheduleAtFixedRate(
            { cleanupExpired() },
            sessionTimeoutMillis,
            sessionTimeoutMillis,
            TimeUnit.MILLISECONDS
        )
    }
    override fun afterConnectionEstablished(session: WebSocketSession) {
        logger.info("Upstream connected: ${session.id}")
        sessions[session.id] = WebSocketProxySession(session)
        onUpstreamOpen(sessions[session.id]!!)
    }
    override fun handleMessage(session: WebSocketSession, message: WebSocketMessage<*>) {
        val proxy = sessions[session.id] ?: return
        if (!proxy.authorized) {
            val ok = onUpstreamFirstMessage(proxy, message)
            if (!ok) {
                onUpstreamFirstMessageIsNull(proxy)
                closeSession(session.id)
                return
            }
            proxy.authorized = true
            onAuthSuccess(proxy)
            connectDownstream(session.id)
            downstreamContexts[session.id]?.pending?.offer(clone(message))
            return
        }
        val ctx = downstreamContexts[session.id] ?: return
        if (!ctx.downConnected.get()) {
            ctx.pending.offer(clone(message))
        } else {
            ctx.sendToDownstream(transformUpstream(message))
        }
    }
    override fun afterConnectionClosed(session: WebSocketSession, status: CloseStatus) {
        logger.info("Upstream closed: ${session.id}")
        closeSession(session.id)
    }
    /**
     * 向所有上游会话发送心跳，维持长连接
     */
    fun sendHeartbeat() {
        val ping = PingMessage()
        sessions.values.forEach {
            try {
                it.session.sendMessage(ping)
            } catch (_: Exception) {
                // 忽略发送失败
            }
        }
    }
    // ---------- 可覆盖钩子 ----------
    /** 上游连接建立后回调 */
    protected open fun onUpstreamOpen(proxy: WebSocketProxySession) = Unit
    /**
     * 上游首条消息处理并授权
     * @return true 表示通过，false 则触发授权失败
     */
    protected abstract fun onUpstreamFirstMessage(
        proxy: WebSocketProxySession,
        message: WebSocketMessage<*>
    ): Boolean
    /** 授权失败发送给上游的消息 */
    protected open fun onUpstreamFirstMessageIsNull(proxy: WebSocketProxySession) {
        val err = mapOf("finish" to true, "error" to "身份认证失败")
        proxy.session.sendMessage(TextMessage(objectMapper.writeValueAsString(err)))
    }
    /** 授权成功后回调 */
    protected open fun onAuthSuccess(proxy: WebSocketProxySession) = Unit
    /** 根据上游会话获取下游 URI */
    protected abstract fun downstreamUri(proxy: WebSocketProxySession): String
    /** 上游→下游 消息转换 */
    protected abstract fun transformUpstream(message: WebSocketMessage<*>): WebSocketMessage<*>
    /** 下游→上游 消息转换 */
    protected abstract fun transformDownstream(message: WebSocketMessage<*>): WebSocketMessage<*>
    /** 会话关闭后回调 */
    protected open fun onSessionClosed(proxy: WebSocketProxySession) = Unit
    // ---------- 内部逻辑 ----------
    /**
     * 建立下游连接，并将后续消息路由到 DownstreamContext
     */
    private fun connectDownstream(sessionId: String) {
        val proxy = sessions[sessionId]!!
        val ctx = DownstreamContext(proxy)
        downstreamContexts[sessionId] = ctx
        client.execute(object : AbstractWebSocketHandler() {
            override fun afterConnectionEstablished(down: WebSocketSession) {
                logger.info("Downstream connected for: $sessionId")
                ctx.downConnected.set(true)
                ctx.downstream = down
                while (true) {
                    val msg = ctx.pending.poll() ?: break
                    ctx.sendToDownstream(transformUpstream(msg))
                }
            }
            override fun handleMessage(down: WebSocketSession, msg: WebSocketMessage<*>) {
                proxy.session.sendMessage(transformDownstream(msg))
            }
            override fun afterConnectionClosed(down: WebSocketSession, status: CloseStatus) {
                logger.warn("Downstream closed early: ${status.code}")
                closeSession(sessionId)
            }
        }, downstreamUri(proxy))
    }
    /** 关闭并清理指定会话 */
    private fun closeSession(sessionId: String) {
        sessions.remove(sessionId)?.also { onSessionClosed(it) }
        downstreamContexts.remove(sessionId)?.closeAll()
    }
    /** 清理超时会话 */
    private fun cleanupExpired() {
        val now = System.currentTimeMillis()
        sessions.entries
            .filter { now - it.value.lastHeartbeat > sessionTimeoutMillis }
            .forEach { closeSession(it.key) }
    }
    /** 克隆消息以避免并发问题 */
    private fun clone(msg: WebSocketMessage<*>): WebSocketMessage<*> = when (msg) {
        is TextMessage -> TextMessage(msg.payload)
        is BinaryMessage -> BinaryMessage(msg.payload.asReadOnlyBuffer())
        else -> msg
    }
    /**
     * 管理下游消息发送及队列
     */
    private class DownstreamContext(proxy: WebSocketProxySession) {
        @Volatile var downstream: WebSocketSession? = null
        val downConnected = AtomicBoolean(false)
        val pending = ConcurrentLinkedQueue<WebSocketMessage<*>>()
        private val executor: ExecutorService = ThreadPoolExecutor(
            4, 16, 60, TimeUnit.SECONDS,
            LinkedBlockingQueue(1000),
            NamedThreadFactory("proxy-send-${proxy.session.id}")
        )
        /** 将消息异步发送到下游 */
        fun sendToDownstream(msg: WebSocketMessage<*>) {
            executor.execute {
                try {
                    downstream?.sendMessage(msg)
                } catch (e: Exception) {
                    LoggerFactory.getLogger("DownstreamLogger").error("Send downstream failed", e)
                }
            }
        }
        /** 关闭下游并清理资源 */
        fun closeAll() {
            try {
                downstream?.close()
            } catch (_: Exception) {
            }
            executor.shutdownNow()
            pending.clear()
        }
    }
    /** 为线程池生成可读性线程名 */
    private class NamedThreadFactory(prefix: String) : ThreadFactory {
        private val cnt = AtomicInteger(1)
        private val name = "${prefix}-${cnt.getAndIncrement()}"
        override fun newThread(r: Runnable) = Thread(r, name)
    }
}