package com.opendance.studio.net

import android.os.Build
import android.util.Log
import com.opendance.studio.BuildConfig
import com.opendance.studio.model.KeypointFrame
import kotlinx.coroutines.CoroutineScope
import kotlinx.coroutines.Dispatchers
import kotlinx.coroutines.Job
import kotlinx.coroutines.delay
import kotlinx.coroutines.flow.MutableStateFlow
import kotlinx.coroutines.flow.StateFlow
import kotlinx.coroutines.flow.asStateFlow
import kotlinx.coroutines.launch
import okhttp3.OkHttpClient
import okhttp3.Request
import okhttp3.Response
import okhttp3.WebSocket
import okhttp3.WebSocketListener
import org.json.JSONObject
import java.util.concurrent.TimeUnit
import java.util.concurrent.atomic.AtomicInteger
import kotlin.math.min

/**
 * T-AND-005：WebSocket 推流客户端（本端为**发送方**）
 *
 * 契约 contracts/websocket-keypoints.md：
 *   - 端点 ws://<windows-host>:8765/pose，子协议 opendance.pose.v1
 *   - 二进制帧 = 239 字节关键点（见 [KeypointFrame]）
 *   - 文本控制帧 hello / config / stats / bye
 *
 * 关键设计：
 *   - 二进制帧走 `send(ByteString)`，控制帧走 `send(text)`
 *   - 自动重连采用指数退避（1s → 2s → 4s … 上限 15s），连接成功即重置
 *   - `sendFrame` 在非连接态直接丢弃并计数，绝不阻塞相机线程
 *   - 统计信息（发送帧数、丢弃帧数、FPS、RTT）由本类维护，供运行页展示
 */
class PoseWebSocketClient(
    private val scope: CoroutineScope,
    private val subprotocol: String = BuildConfig.WS_SUBPROTOCOL,
) {

    enum class State { IDLE, CONNECTING, CONNECTED, RECONNECTING, CLOSED }

    data class Stats(
        val state: State = State.IDLE,
        val url: String = "",
        val sentFrames: Int = 0,
        val droppedFrames: Int = 0,
        val lastFrameId: Int = -1,
        val sendFps: Float = 0f,
        val roundTripMs: Long = -1,
        val reconnectAttempts: Int = 0,
        val lastError: String = "",
        val serverSessionId: String = "",
    )

    private val client: OkHttpClient = OkHttpClient.Builder()
        .connectTimeout(CONNECT_TIMEOUT_SECONDS, TimeUnit.SECONDS)
        .readTimeout(0, TimeUnit.MILLISECONDS)   // WebSocket 长连接不设读超时
        .writeTimeout(WRITE_TIMEOUT_SECONDS, TimeUnit.SECONDS)
        .pingInterval(PING_INTERVAL_SECONDS, TimeUnit.SECONDS)
        .retryOnConnectionFailure(true)
        .build()

    private val _stats = MutableStateFlow(Stats())
    val stats: StateFlow<Stats> = _stats.asStateFlow()

    private var webSocket: WebSocket? = null
    private var reconnectJob: Job? = null
    private var url: String = ""
    private var manualClose = false

    private val sentCounter = AtomicInteger(0)
    private val droppedCounter = AtomicInteger(0)
    private val reconnectCounter = AtomicInteger(0)
    private val frameIdCounter = AtomicInteger(0)

    /** 首帧时间戳与计数，用于计算发送 FPS */
    private var fpsWindowStart = 0L
    private var fpsWindowFrames = 0

    private val deviceLabel: String =
        "${Build.MANUFACTURER} ${Build.MODEL}".trim().ifBlank { "Android" }

    // ------------------------------------------------------------------ 连接

    fun connect(targetUrl: String, autoReconnect: Boolean = true) {
        url = targetUrl
        manualClose = false
        this.autoReconnect = autoReconnect
        reconnectCounter.set(0)
        openSocket(initial = true)
    }

    private var autoReconnect = true

    private fun openSocket(initial: Boolean) {
        closeSocketQuietly()

        val state = if (initial) State.CONNECTING else State.RECONNECTING
        update(state = state, error = "")

        val request = Request.Builder()
            .url(url)
            .header("Sec-WebSocket-Protocol", subprotocol)
            .build()

        Log.i(TAG, "连接 $url（子协议 $subprotocol）")
        webSocket = client.newWebSocket(request, listener)
    }

    private val listener = object : WebSocketListener() {

        override fun onOpen(webSocket: WebSocket, response: Response) {
            Log.i(TAG, "已连接：${response.code}")
            reconnectCounter.set(0)
            fpsWindowStart = 0L
            fpsWindowFrames = 0
            update(state = State.CONNECTED, error = "")
            sendHello()
        }

        override fun onMessage(webSocket: WebSocket, text: String) {
            handleControlMessage(text)
        }

        override fun onClosing(webSocket: WebSocket, code: Int, reason: String) {
            Log.i(TAG, "服务端关闭：$code $reason")
            webSocket.close(1000, null)
        }

        override fun onClosed(webSocket: WebSocket, code: Int, reason: String) {
            Log.i(TAG, "连接已关闭：$code $reason")
            update(state = State.CLOSED)
            scheduleReconnect()
        }

        override fun onFailure(webSocket: WebSocket, t: Throwable, response: Response?) {
            val reason = t.message ?: t.javaClass.simpleName
            Log.w(TAG, "连接失败：$reason")
            update(state = State.CLOSED, error = reason)
            scheduleReconnect()
        }
    }

    private fun scheduleReconnect() {
        if (manualClose || !autoReconnect) return

        reconnectJob?.cancel()
        val attempt = reconnectCounter.incrementAndGet()
        val delayMs = min(BASE_RECONNECT_MS shl (attempt - 1), MAX_RECONNECT_MS.toInt()).toLong()
        update(state = State.RECONNECTING, attempts = attempt)

        reconnectJob = scope.launch {
            delay(delayMs)
            if (!manualClose && autoReconnect) {
                Log.i(TAG, "第 $attempt 次重连（等待 ${delayMs}ms）")
                openSocket(initial = false)
            }
        }
    }

    fun disconnect() {
        manualClose = true
        reconnectJob?.cancel()
        reconnectJob = null
        sendBye("user_stop")
        closeSocketQuietly()
        update(state = State.CLOSED)
    }

    private fun closeSocketQuietly() {
        webSocket?.let {
            it.cancel()
        }
        webSocket = null
    }

    // ------------------------------------------------------------------ 发送

    /**
     * 发送一帧关键点。非连接态直接丢弃并计入 droppedFrames —— 绝不能阻塞相机线程。
     * @return 实际是否发出
     */
    fun sendFrame(frame: KeypointFrame.PoseFrame): Boolean {
        val socket = webSocket ?: return drop()
        if (_stats.value.state != State.CONNECTED) return drop()

        val frameId = frameIdCounter.getAndIncrement()
        val bytes = KeypointFrame.encode(
            frameId = frameId,
            xs = frame.x,
            ys = frame.y,
            zs = frame.z,
            confidence = frame.confidence,
            flags = frame.flags,
        )

        val ok = socket.send(okio.ByteString.of(*bytes))
        if (!ok) return drop()

        sentCounter.incrementAndGet()
        fpsWindowFrames++
        val now = System.currentTimeMillis()
        if (fpsWindowStart == 0L) fpsWindowStart = now
        val elapsed = now - fpsWindowStart
        val fps = if (elapsed >= FPS_WINDOW_MS) {
            val value = fpsWindowFrames * 1000f / elapsed
            fpsWindowStart = now
            fpsWindowFrames = 0
            value
        } else {
            _stats.value.sendFps
        }

        _stats.value = _stats.value.copy(
            sentFrames = sentCounter.get(),
            droppedFrames = droppedCounter.get(),
            lastFrameId = frameId,
            sendFps = fps,
        )
        return true
    }

    private fun drop(): Boolean {
        droppedCounter.incrementAndGet()
        _stats.value = _stats.value.copy(droppedFrames = droppedCounter.get())
        return false
    }

    fun sendHello() {
        sendJson(
            JSONObject().apply {
                put("type", "hello")
                put("protocol", subprotocol)
                put("device", deviceLabel)
                put("appVersion", BuildConfig.VERSION_NAME)
                put("cameraFacing", "front")
                put("targetFps", TARGET_FPS_FIELD)
                put("resolution", RESOLUTION_FIELD)
            },
        )
    }

    fun sendStats(fps: Float, dropped: Int, latencyMs: Long, confidenceAvg: Float) {
        sendJson(
            JSONObject().apply {
                put("type", "stats")
                put("fps", fps.toDouble())
                put("dropped", dropped)
                put("latencyMs", latencyMs)
                put("confidenceAvg", confidenceAvg.toDouble())
            },
        )
    }

    fun sendBye(reason: String) {
        sendJson(
            JSONObject().apply {
                put("type", "bye")
                put("reason", reason)
            },
        )
    }

    private fun sendJson(json: JSONObject) {
        val socket = webSocket ?: return
        if (_stats.value.state != State.CONNECTED) return
        socket.send(json.toString())
    }

    private fun handleControlMessage(text: String) {
        try {
            val json = JSONObject(text)
            when (json.optString("type")) {
                "welcome" -> {
                    _stats.value = _stats.value.copy(
                        serverSessionId = json.optString("sessionId", ""),
                        roundTripMs = 0,
                    )
                }
                "config" -> onRemoteConfig?.invoke(json)
                "bye" -> Log.i(TAG, "服务端要求断开：${json.optString("reason")}")
            }
        } catch (e: Exception) {
            Log.w(TAG, "控制帧解析失败：${e.message}")
        }
    }

    /** 服务端下发 config 时回调（契约 §3.3） */
    var onRemoteConfig: ((JSONObject) -> Unit)? = null

    private fun update(
        state: State,
        error: String = _stats.value.lastError,
        attempts: Int = _stats.value.reconnectAttempts,
    ) {
        _stats.value = _stats.value.copy(
            state = state,
            url = url,
            lastError = error,
            reconnectAttempts = attempts,
            sentFrames = sentCounter.get(),
            droppedFrames = droppedCounter.get(),
        )
    }

    fun release() {
        disconnect()
        client.dispatcher.executorService.shutdown()
    }

    companion object {
        private const val TAG = "PoseWebSocket"
        private const val CONNECT_TIMEOUT_SECONDS = 8L
        private const val WRITE_TIMEOUT_SECONDS = 5L
        private const val PING_INTERVAL_SECONDS = 15L
        private const val BASE_RECONNECT_MS = 1000
        private const val MAX_RECONNECT_MS = 15000L
        private const val FPS_WINDOW_MS = 1000L

        /** hello 里的字段由设置填充，运行时由调用方覆盖 */
        @Volatile var TARGET_FPS_FIELD: Int = 30
        @Volatile var RESOLUTION_FIELD: String = "640x480"
    }
}