package app.dudu.net

import okhttp3.OkHttpClient
import okhttp3.Request
import okhttp3.Response
import okhttp3.WebSocket
import okhttp3.WebSocketListener
import okio.ByteString.Companion.toByteString
import org.json.JSONObject
import java.util.concurrent.TimeUnit

/**
 * WS 客户端，协议对齐 server/pkg/protocol/ws.go（WSS /v1/asr/stream）：
 *  - 握手 header：Authorization: Bearer <jwt>、X-Device-ID
 *  - 上行文本帧：{"type":"start","session_id":...,"sample_rate":16000,"format":"pcm16"} / stop / cancel
 *  - 上行二进制帧：PCM 16k/16bit/mono，100ms = 3200B
 *  - 下行文本帧：partial / final / usage / error
 *
 * MVP：每次按住-松开新建一条连接（简单可靠）。
 * TODO(latency): 改为常驻连接 + ping/pong 保活，按键零握手（见 backend-architecture 第四章）。
 */
class AsrSocket(
    private val wsUrl: String,
    private val token: String,
    private val deviceId: String,
    private val listener: Listener,
) {
    interface Listener {
        fun onOpen()
        fun onPartial(sessionId: String, text: String)
        fun onFinal(sessionId: String, text: String)
        fun onUsage(sessionSeconds: Int, balanceSeconds: Long, trialRemaining: Int)
        fun onServerError(code: String, message: String)
        fun onClosed()
        fun onFailure(t: Throwable)
    }

    private var ws: WebSocket? = null

    fun connect() {
        val req = Request.Builder()
            .url(wsUrl)
            .header("Authorization", "Bearer $token")
            .header("X-Device-ID", deviceId)
            .build()
        ws = client.newWebSocket(req, object : WebSocketListener() {
            override fun onOpen(webSocket: WebSocket, response: Response) = listener.onOpen()

            override fun onMessage(webSocket: WebSocket, text: String) {
                val msg = runCatching { JSONObject(text) }.getOrNull() ?: return
                val sid = msg.optString("session_id")
                when (msg.optString("type")) {
                    "partial" -> listener.onPartial(sid, msg.optString("text"))
                    "final" -> listener.onFinal(sid, msg.optString("text"))
                    "usage" -> listener.onUsage(
                        msg.optInt("session_seconds"),
                        msg.optLong("balance_seconds"),
                        msg.optInt("trial_remaining"),
                    )
                    "error" -> listener.onServerError(msg.optString("code"), msg.optString("message"))
                }
            }

            override fun onClosed(webSocket: WebSocket, code: Int, reason: String) = listener.onClosed()

            override fun onFailure(webSocket: WebSocket, t: Throwable, response: Response?) =
                listener.onFailure(t)
        })
    }

    fun sendStart(sessionId: String) {
        val json = JSONObject()
            .put("type", "start")
            .put("session_id", sessionId)
            .put("sample_rate", 16000)
            .put("format", "pcm16")
        ws?.send(json.toString())
    }

    fun sendAudio(frame: ByteArray) {
        ws?.send(frame.toByteString(0, frame.size))
    }

    fun sendStop(sessionId: String) {
        ws?.send(JSONObject().put("type", "stop").put("session_id", sessionId).toString())
    }

    fun sendCancel(sessionId: String) {
        ws?.send(JSONObject().put("type", "cancel").put("session_id", sessionId).toString())
    }

    fun close() {
        ws?.close(1000, null)
        ws = null
    }

    companion object {
        // 进程级单例 client：连接池/线程池复用
        private val client: OkHttpClient by lazy {
            OkHttpClient.Builder()
                .connectTimeout(5, TimeUnit.SECONDS)
                .pingInterval(20, TimeUnit.SECONDS)
                .build()
        }
    }
}
