package com.opendance.studio.model

import java.nio.ByteBuffer
import java.nio.ByteOrder

/**
 * 239 字节关键点帧编解码 —— 契约 contracts/websocket-keypoints.md §2（冻结）
 *
 * 帧布局（**小端序**）：
 * ```
 * 偏移  长度  字段
 * 0     2     magic    u16 = 0x4F44 ('OD')
 * 2     1     version  u8  = 1
 * 3     1     flags    u8  bit0 镜像 | bit1 含低置信度点 | bit2 上半身模式
 * 4     4     frame_id u32 LE，从 0 起单调递增
 * 8     231   keypoints 33 × 7B：s16 x, s16 y, s16 z, u8 conf
 * 合计  239
 * ```
 *
 * 坐标语义（contracts/skeleton-topology.md §4）：
 * x,y ∈ [0,1] 归一化图像坐标（原点左上，y 向下）；z 与 x 同尺度，
 * 负值朝向相机。归一化系数统一 **32767**。
 *
 * 本实现与 opendance-web 的 TypeScript 参考实现、opendance-windows 的 GDScript
 * 实现逐位对齐，并由单元测试用同一份金标准字节向量交叉校验。
 */
object KeypointFrame {

    const val MAGIC = 0x4F44
    const val PROTOCOL_VERSION = 1
    const val BINARY_SIZE = 239
    const val KEYPOINT_COUNT = 33
    const val HEADER_SIZE = 8
    const val POINT_BYTES = 7
    const val SCALE = 32767.0

    const val FLAG_MIRRORED = 1 shl 0
    const val FLAG_LOW_CONFIDENCE = 1 shl 1
    const val FLAG_UPPER_BODY_ONLY = 1 shl 2

    /** 低于该置信度的点坐标无意义，接收端必须插值 */
    const val CONFIDENCE_THRESHOLD = 0.5f

    /** 一帧姿态数据（并行数组，避免 30 FPS 循环里的对象分配） */
    class PoseFrame(
        val frameId: Int,
        val flags: Int,
        val x: FloatArray,
        val y: FloatArray,
        val z: FloatArray,
        val confidence: FloatArray,
    ) {
        init {
            require(x.size == KEYPOINT_COUNT) { "x 必须为 $KEYPOINT_COUNT 项" }
            require(y.size == KEYPOINT_COUNT) { "y 必须为 $KEYPOINT_COUNT 项" }
            require(z.size == KEYPOINT_COUNT) { "z 必须为 $KEYPOINT_COUNT 项" }
            require(confidence.size == KEYPOINT_COUNT) { "confidence 必须为 $KEYPOINT_COUNT 项" }
        }

        val isMirrored: Boolean get() = (flags and FLAG_MIRRORED) != 0
        val isUpperBodyOnly: Boolean get() = (flags and FLAG_UPPER_BODY_ONLY) != 0
        val hasLowConfidence: Boolean get() = (flags and FLAG_LOW_CONFIDENCE) != 0

        /** 该点是否可用于评分 */
        fun isReliable(index: Int): Boolean = confidence[index] >= CONFIDENCE_THRESHOLD

        /** 平均置信度，用于运行页质量指示 */
        fun averageConfidence(): Float {
            var total = 0f
            for (c in confidence) total += c
            return total / confidence.size
        }

        /** 已镜像的帧需翻转 x 才能与对端坐标系一致 */
        fun xUnmirrored(index: Int): Float = if (isMirrored) 1f - x[index] else x[index]
    }

    /**
     * 校验帧字节流；合法返回 null，否则返回可读原因。
     */
    fun validate(bytes: ByteArray): String? {
        if (bytes.size != BINARY_SIZE) {
            return "帧长度必须为 $BINARY_SIZE 字节，实际 ${bytes.size}"
        }
        val buffer = ByteBuffer.wrap(bytes).order(ByteOrder.LITTLE_ENDIAN)
        val magic = buffer.getShort(0).toInt() and 0xFFFF
        if (magic != MAGIC) {
            return "magic 校验失败：期望 0x%04X，实际 0x%04X".format(MAGIC, magic)
        }
        val version = buffer.get(2).toInt() and 0xFF
        if (version != PROTOCOL_VERSION) {
            return "不支持的协议版本 $version（本端支持 $PROTOCOL_VERSION）"
        }
        return null
    }

    /**
     * 编码为 239 字节帧 —— 本端（Android）的主要职责。
     *
     * @param xs 归一化 x，[0,1]
     * @param ys 归一化 y，[0,1]
     * @param zs 深度，约 [-1,1]，负值朝相机
     * @param confidence [0,1]
     * @param flags 见 FLAG_* 常量
     */
    fun encode(
        frameId: Int,
        xs: FloatArray,
        ys: FloatArray,
        zs: FloatArray,
        confidence: FloatArray,
        flags: Int = 0,
    ): ByteArray {
        require(xs.size == KEYPOINT_COUNT) { "xs 必须为 $KEYPOINT_COUNT 项" }
        require(ys.size == KEYPOINT_COUNT) { "ys 必须为 $KEYPOINT_COUNT 项" }
        require(zs.size == KEYPOINT_COUNT) { "zs 必须为 $KEYPOINT_COUNT 项" }
        require(confidence.size == KEYPOINT_COUNT) { "confidence 必须为 $KEYPOINT_COUNT 项" }

        val buffer = ByteBuffer.allocate(BINARY_SIZE).order(ByteOrder.LITTLE_ENDIAN)
        buffer.putShort(0, MAGIC.toShort())
        buffer.put(2, PROTOCOL_VERSION.toByte())
        buffer.put(3, (flags and 0xFF).toByte())
        buffer.putInt(4, frameId)

        for (i in 0 until KEYPOINT_COUNT) {
            val offset = HEADER_SIZE + i * POINT_BYTES
            buffer.putShort(offset, toFixed16(xs[i]))
            buffer.putShort(offset + 2, toFixed16(ys[i]))
            buffer.putShort(offset + 4, toFixed16(zs[i]))
            buffer.put(offset + 6, toConfidenceByte(confidence[i]))
        }
        return buffer.array()
    }

    /** 解码 239 字节帧；非法输入抛 IllegalArgumentException */
    fun decode(bytes: ByteArray): PoseFrame {
        validate(bytes)?.let { throw IllegalArgumentException(it) }

        val buffer = ByteBuffer.wrap(bytes).order(ByteOrder.LITTLE_ENDIAN)
        val frameId = buffer.getInt(4)
        val flags = buffer.get(3).toInt() and 0xFF

        val xs = FloatArray(KEYPOINT_COUNT)
        val ys = FloatArray(KEYPOINT_COUNT)
        val zs = FloatArray(KEYPOINT_COUNT)
        val confidence = FloatArray(KEYPOINT_COUNT)

        for (i in 0 until KEYPOINT_COUNT) {
            val offset = HEADER_SIZE + i * POINT_BYTES
            xs[i] = buffer.getShort(offset) / SCALE.toFloat()
            ys[i] = buffer.getShort(offset + 2) / SCALE.toFloat()
            zs[i] = buffer.getShort(offset + 4) / SCALE.toFloat()
            confidence[i] = (buffer.get(offset + 6).toInt() and 0xFF) / 255f
        }
        return PoseFrame(frameId, flags, xs, ys, zs, confidence)
    }

    /** 依据置信度数组计算是否应置 FLAG_LOW_CONFIDENCE */
    fun computeLowConfidence(confidence: FloatArray): Boolean =
        confidence.any { it < CONFIDENCE_THRESHOLD }

    fun withLowConfidenceFlag(flags: Int, confidence: FloatArray): Int =
        if (computeLowConfidence(confidence)) flags or FLAG_LOW_CONFIDENCE else flags and FLAG_LOW_CONFIDENCE.inv()

    private fun toFixed16(value: Float): Short {
        val scaled = Math.round(value * SCALE)
        return scaled.coerceIn(-32768L, 32767L).toShort()
    }

    private fun toConfidenceByte(value: Float): Byte =
        Math.round(value.coerceIn(0f, 1f) * 255f).coerceIn(0, 255).toByte()
}