path: root/core/src/commonMain/kotlin/coffee/liz/rookery/bt/GanProtocol.kt

// Protocol details ported from gan-web-bluetooth by Andy Fedotov, MIT License
// (https://github.com/afedotov/gan-web-bluetooth).
package coffee.liz.rookery.bt

import coffee.liz.rookery.bt.SmartCubeEvent.Move

enum class GanCommand { REQUEST_HARDWARE, REQUEST_FACELETS, REQUEST_BATTERY, REQUEST_RESET }

/** What a protocol driver may do back to the cube while handling an event (move-history requests, disconnect). */
interface GanRawConnection {
    /** Encrypts and writes a plaintext command message. */
    suspend fun sendCommandMessage(message: ByteArray)
    suspend fun disconnect()
}

/** Encodes commands and decodes decrypted state notifications for one GAN protocol generation. */
interface GanProtocolDriver {
    fun createCommandMessage(command: GanCommand): ByteArray?
    suspend fun handleStateEvent(conn: GanRawConnection, message: ByteArray): List<SmartCubeEvent>

    /** Called every ~150 ms so drivers can retry or give up on things the cube never answered. */
    suspend fun tick(conn: GanRawConnection): List<SmartCubeEvent> = emptyList()
}

private const val FACES = "URFDLB"
private fun moveName(face: Int, direction: Int) = FACES[face].toString() + (if (direction == 1) "'" else "")

private fun msgOf(size: Int, vararg head: Int) = ByteArray(size).also { b -> head.forEachIndexed { i, v -> b[i] = v.toByte() } }

private fun signed15(q: Int): Double = (1 - (q shr 15) * 2) * (q and 0x7FFF) / 0x7FFF.toDouble()
private fun signed3(v: Int): Int = (1 - (v shr 3) * 2) * (v and 0x7)

private fun readState(r: GanBitReader, cpAt: Int, coAt: Int, epAt: Int, eoAt: Int): CubeState = GanFacelets.complete(
    cp7 = List(7) { r.bits(cpAt + it * 3, 3) },
    co7 = List(7) { r.bits(coAt + it * 2, 2) },
    ep11 = List(11) { r.bits(epAt + it * 4, 4) },
    eo11 = List(11) { r.bits(eoAt + it, 1) },
)

private fun facelets(state: CubeState, serial: Int) = SmartCubeEvent.Facelets(GanFacelets.toKociemba(state), serial, state)

private fun gyro(r: GanBitReader, at: Int) = SmartCubeEvent.Gyro(
    Quaternion(
        x = signed15(r.bits(at + 16, 16)),
        y = signed15(r.bits(at + 32, 16)),
        z = signed15(r.bits(at + 48, 16)),
        w = signed15(r.bits(at, 16)),
    ),
    AngularVelocity(signed3(r.bits(at + 64, 4)), signed3(r.bits(at + 68, 4)), signed3(r.bits(at + 72, 4))),
)

/**
 * GAN Gen2: GAN Mini ui FreePlay, GAN12 ui (FreePlay), GAN356 i Carry (S), GAN356 i 3, Monster Go 3Ai, MoYu AI 2023.
 * 20-byte messages; event type in the first nibble.
 */
class GanGen2ProtocolDriver(private val clock: () -> Long) : GanProtocolDriver {
    private var lastSerial = -1
    private var lastMoveTimestamp = 0L
    private var cubeTimestamp = 0L

    override fun createCommandMessage(command: GanCommand): ByteArray = when (command) {
        GanCommand.REQUEST_FACELETS -> msgOf(20, 0x04)
        GanCommand.REQUEST_HARDWARE -> msgOf(20, 0x05)
        GanCommand.REQUEST_BATTERY -> msgOf(20, 0x09)
        GanCommand.REQUEST_RESET -> msgOf(20, 0x0A, 0x05, 0x39, 0x77, 0x00, 0x00, 0x01, 0x23, 0x45, 0x67, 0x89, 0xAB)
    }

    override suspend fun handleStateEvent(conn: GanRawConnection, message: ByteArray): List<SmartCubeEvent> {
        val timestamp = clock()
        val events = mutableListOf<SmartCubeEvent>()
        val r = GanBitReader(message)
        when (r.bits(0, 4)) {
            0x01 -> events += gyro(r, 4)
            0x02 -> if (lastSerial != -1) { // moves are accepted only after the first facelets event
                val serial = r.bits(4, 8)
                val diff = minOf((serial - lastSerial) and 0xFF, 7)
                lastSerial = serial
                if (diff > 0) {
                    for (i in diff - 1 downTo 0) {
                        val face = r.bits(12 + 5 * i, 4)
                        val direction = r.bits(16 + 5 * i, 1)
                        var elapsed = r.bits(47 + 16 * i, 16).toLong()
                        if (elapsed == 0L) elapsed = timestamp - lastMoveTimestamp // 16-bit cube timer overflow
                        cubeTimestamp += elapsed
                        if (face < FACES.length) {
                            events += Move(moveName(face, direction), cubeTimestamp, timestamp, (serial - i) and 0xFF, recovered = i != 0)
                        }
                    }
                    lastMoveTimestamp = timestamp
                }
            }
            0x04 -> {
                val serial = r.bits(4, 8)
                if (lastSerial == -1) lastSerial = serial
                events += facelets(readState(r, 12, 33, 47, 91), serial)
            }
            0x05 -> {
                val hwMajor = r.bits(8, 8); val hwMinor = r.bits(16, 8)
                val swMajor = r.bits(24, 8); val swMinor = r.bits(32, 8)
                val gyroSupported = r.bits(104, 1) == 1
                val name = (0 until 8).map { r.bits(it * 8 + 40, 8).toChar() }.joinToString("")
                events += SmartCubeEvent.Hardware(name.trimEnd('\u0000'), "$hwMajor.$hwMinor", "$swMajor.$swMinor", null, gyroSupported)
            }
            0x09 -> events += SmartCubeEvent.Battery(minOf(r.bits(8, 8), 100))
            0x0D -> conn.disconnect()
        }
        return events
    }
}

/**
 * Shared move FIFO with gap detection and move-history recovery used by Gen3 and Gen4:
 * live moves carry a serial; when one is skipped, the driver asks the cube for its move history
 * and holds later moves back until the gap is filled.
 */
abstract class GanBufferedProtocolDriver(protected val clock: () -> Long) : GanProtocolDriver {
    protected var serial = -1
    protected var lastSerial = -1
    protected var lastLocalTimestamp: Long? = null
    private val moveBuffer = ArrayDeque<Move>()

    protected abstract fun moveHistoryMessage(serial: Int, count: Int): ByteArray

    private suspend fun requestMoveHistory(conn: GanRawConnection, serialIn: Int, countIn: Int) {
        // History data is byte-aligned and always starts at an odd serial: request an odd-aligned, even-sized window,
        // and never across the 255 -> 0 edge (firmware spoofs those as 'D').
        var s = serialIn
        var count = countIn
        if (s % 2 == 0) s = (s - 1) and 0xFF
        if (count % 2 == 1) count++
        count = minOf(count, s + 1)
        try {
            conn.sendCommandMessage(moveHistoryMessage(s, count))
        } catch (e: Exception) {
            if (e is kotlinx.coroutines.CancellationException) throw e
            // Ignored: retried on the next move event.
        }
    }

    /** When the buffer started waiting on a gap (null = not stalled), and when we last asked for history. */
    private var stalledSince: Long? = null
    private var lastHistoryRequest = 0L

    protected suspend fun evictMoveBuffer(conn: GanRawConnection?): List<SmartCubeEvent> {
        val evicted = mutableListOf<SmartCubeEvent>()
        while (moveBuffer.isNotEmpty()) {
            val head = moveBuffer.first()
            val diff = if (lastSerial == -1) 1 else (head.serial - lastSerial) and 0xFF
            if (diff > 1) {
                if (stalledSince == null) stalledSince = clock()
                if (conn != null) { requestMoveHistory(conn, head.serial, diff); lastHistoryRequest = clock() }
                break
            } else {
                evicted += moveBuffer.removeFirst()
                lastSerial = head.serial
            }
        }
        if (moveBuffer.isEmpty()) stalledSince = null
        // A big backlog means history isn't coming: stop waiting rather than hold moves hostage.
        if (moveBuffer.size > MAX_BUFFERED) evicted += giveUp(conn)
        return evicted
    }

    /**
     * Release everything held behind a gap, in order. The first released move is marked [Move.afterGap]
     * so consumers know moves are missing before it; we also ask for facelets so they can resync.
     */
    private suspend fun giveUp(conn: GanRawConnection?): List<SmartCubeEvent> {
        val released = moveBuffer.mapIndexed { i, m -> if (i == 0) m.copy(afterGap = true) else m }
        moveBuffer.clear()
        released.lastOrNull()?.let { lastSerial = it.serial }
        stalledSince = null
        if (conn != null) createCommandMessage(GanCommand.REQUEST_FACELETS)?.let { msg ->
            try { conn.sendCommandMessage(msg) } catch (e: Exception) { if (e is kotlinx.coroutines.CancellationException) throw e }
        }
        return released
    }

    /** Retry unanswered history requests, and give up on a gap that history can't fill. */
    override suspend fun tick(conn: GanRawConnection): List<SmartCubeEvent> {
        val since = stalledSince ?: return emptyList()
        val now = clock()
        if (now - since >= GIVE_UP_MS) return giveUp(conn)
        if (now - lastHistoryRequest >= RETRY_MS) {
            val head = moveBuffer.firstOrNull() ?: return emptyList()
            requestMoveHistory(conn, head.serial, (head.serial - lastSerial) and 0xFF)
            lastHistoryRequest = now
        }
        return emptyList()
    }

    private fun isSerialInRange(start: Int, end: Int, serial: Int, closedStart: Boolean = false, closedEnd: Boolean = false) =
        ((end - start) and 0xFF) >= ((serial - start) and 0xFF) &&
            (closedStart || ((start - serial) and 0xFF) > 0) &&
            (closedEnd || ((end - serial) and 0xFF) > 0)

    private fun injectMissedMove(move: Move) {
        if (moveBuffer.isNotEmpty()) {
            val head = moveBuffer.first()
            if (moveBuffer.any { it.serial == move.serial }) return
            if (!isSerialInRange(lastSerial, head.serial, move.serial)) return
            // History is injected in reverse order, so only the move right before the head fits.
            if (move.serial == ((head.serial - 1) and 0xFF)) moveBuffer.addFirst(move)
        } else if (isSerialInRange(lastSerial, serial, move.serial, closedEnd = true)) {
            // A lost move recovered via the periodic facelets event, into an empty buffer.
            moveBuffer.addFirst(move)
        }
    }

    private suspend fun checkIfMoveMissed(conn: GanRawConnection) {
        val diff = (serial - lastSerial) and 0xFF
        if (diff > 0 && serial != 0) { // serial 0: firmware bug with facelets event at counter 255
            val startSerial = moveBuffer.firstOrNull()?.serial ?: ((serial + 1) and 0xFF)
            requestMoveHistory(conn, startSerial, diff + 1)
        }
    }

    protected suspend fun onLiveMove(conn: GanRawConnection, timestamp: Long, cubeTimestamp: Long, moveSerial: Int, direction: Int, faceBits: Int): List<SmartCubeEvent> {
        if (lastSerial == -1) return emptyList() // moves are accepted only after the first facelets event
        lastLocalTimestamp = timestamp
        serial = moveSerial
        val face = LIVE_FACE_BITS.indexOf(faceBits)
        if (face >= 0) moveBuffer.addLast(Move(moveName(face, direction), cubeTimestamp, timestamp, moveSerial))
        return evictMoveBuffer(conn)
    }

    protected suspend fun onMoveHistory(r: GanBitReader, timestamp: Long, startSerial: Int, dataLength: Int, movesAt: Int): List<SmartCubeEvent> {
        val count = (dataLength - 1) * 2
        for (i in 0 until count) {
            val face = HISTORY_FACE_BITS.indexOf(r.bits(movesAt + 4 * i, 3))
            val direction = r.bits(movesAt + 3 + 4 * i, 1)
            if (face >= 0) injectMissedMove(Move(moveName(face, direction), null, timestamp, (startSerial - i) and 0xFF, recovered = true))
        }
        return evictMoveBuffer(null)
    }

    protected suspend fun onFaceletsSerial(conn: GanRawConnection, timestamp: Long, faceletsSerial: Int) {
        serial = faceletsSerial
        // The cube sends facelets periodically; use them to detect missed moves, debounced while turning.
        if (lastSerial != -1) {
            val last = lastLocalTimestamp
            if (last != null && timestamp - last > 500) checkIfMoveMissed(conn)
        }
        if (lastSerial == -1) lastSerial = faceletsSerial
    }

    private companion object {
        const val RETRY_MS = 300L
        const val GIVE_UP_MS = 1_500L
        const val MAX_BUFFERED = 16
        val LIVE_FACE_BITS = listOf(2, 32, 8, 1, 16, 4)
        val HISTORY_FACE_BITS = listOf(1, 5, 3, 0, 4, 2)
    }
}

/** GAN Gen3: GAN356 i Carry 2. 16-byte messages starting with magic 0x55, type, length. */
class GanGen3ProtocolDriver(clock: () -> Long) : GanBufferedProtocolDriver(clock) {
    override fun createCommandMessage(command: GanCommand): ByteArray = when (command) {
        GanCommand.REQUEST_FACELETS -> msgOf(16, 0x68, 0x01)
        GanCommand.REQUEST_HARDWARE -> msgOf(16, 0x68, 0x04)
        GanCommand.REQUEST_BATTERY -> msgOf(16, 0x68, 0x07)
        GanCommand.REQUEST_RESET -> msgOf(16, 0x68, 0x05, 0x05, 0x39, 0x77, 0x00, 0x00, 0x01, 0x23, 0x45, 0x67, 0x89, 0xAB)
    }

    override fun moveHistoryMessage(serial: Int, count: Int) = msgOf(16, 0x68, 0x03, serial, 0, count, 0)

    override suspend fun handleStateEvent(conn: GanRawConnection, message: ByteArray): List<SmartCubeEvent> {
        val timestamp = clock()
        val r = GanBitReader(message)
        val magic = r.bits(0, 8)
        val eventType = r.bits(8, 8)
        val dataLength = r.bits(16, 8)
        if (magic != 0x55 || dataLength <= 0) return emptyList()
        return when (eventType) {
            0x01 -> onLiveMove(conn, timestamp, r.u32(24, true), r.bits(56, 16, true) and 0xFF, r.bits(72, 2), r.bits(74, 6))
            0x06 -> onMoveHistory(r, timestamp, r.bits(24, 8), dataLength, 32)
            0x02 -> {
                val s = r.bits(24, 16, true) and 0xFF
                onFaceletsSerial(conn, timestamp, s)
                listOf(facelets(readState(r, 40, 61, 77, 121), s))
            }
            0x07 -> {
                val swMajor = r.bits(72, 4); val swMinor = r.bits(76, 4)
                val hwMajor = r.bits(80, 4); val hwMinor = r.bits(84, 4)
                val name = (0 until 5).map { r.bits(it * 8 + 32, 8).toChar() }.joinToString("")
                listOf(SmartCubeEvent.Hardware(name.trimEnd('\u0000'), "$hwMajor.$hwMinor", "$swMajor.$swMinor", null, false))
            }
            0x10 -> listOf(SmartCubeEvent.Battery(minOf(r.bits(24, 8), 100)))
            0x11 -> { conn.disconnect(); emptyList() }
            else -> emptyList()
        }
    }
}

/** GAN Gen4: GAN12 ui Maglev, GAN14 ui FreePlay. 20-byte messages starting with type, length. */
class GanGen4ProtocolDriver(clock: () -> Long) : GanBufferedProtocolDriver(clock) {
    private val hwInfo = mutableMapOf<Int, String>()

    override fun createCommandMessage(command: GanCommand): ByteArray = when (command) {
        GanCommand.REQUEST_FACELETS -> msgOf(20, 0xDD, 0x04, 0x00, 0xED, 0x00, 0x00)
        GanCommand.REQUEST_HARDWARE -> { hwInfo.clear(); msgOf(20, 0xDF, 0x03, 0x00, 0x00, 0x00) }
        GanCommand.REQUEST_BATTERY -> msgOf(20, 0xDD, 0x04, 0x00, 0xEF, 0x00, 0x00)
        GanCommand.REQUEST_RESET -> msgOf(20, 0xD2, 0x0D, 0x05, 0x39, 0x77, 0x00, 0x00, 0x01, 0x23, 0x45, 0x67, 0x89, 0xAB)
    }

    override fun moveHistoryMessage(serial: Int, count: Int) = msgOf(20, 0xD1, 0x04, serial, 0, count, 0)

    override suspend fun handleStateEvent(conn: GanRawConnection, message: ByteArray): List<SmartCubeEvent> {
        val timestamp = clock()
        val r = GanBitReader(message)
        val eventType = r.bits(0, 8)
        val dataLength = r.bits(8, 8)
        return when (eventType) {
            0x01 -> onLiveMove(conn, timestamp, r.u32(16, true), r.bits(48, 16, true) and 0xFF, r.bits(64, 2), r.bits(66, 6))
            0xD1 -> onMoveHistory(r, timestamp, r.bits(16, 8), dataLength, 24)
            0xED -> {
                val s = r.bits(16, 16, true) and 0xFF
                onFaceletsSerial(conn, timestamp, s)
                listOf(facelets(readState(r, 32, 53, 69, 113), s))
            }
            in 0xFA..0xFE -> {
                when (eventType) {
                    0xFA -> {
                        val year = r.bits(24, 16, true); val month = r.bits(40, 8); val day = r.bits(48, 8)
                        hwInfo[eventType] = "${year.toString().padStart(4, '0')}-${month.toString().padStart(2, '0')}-${day.toString().padStart(2, '0')}"
                    }
                    0xFC -> hwInfo[eventType] = (0 until maxOf(0, minOf(dataLength - 1, (message.size * 8 - 24) / 8)))
                        .map { r.bits(it * 8 + 24, 8).toChar() }.joinToString("")
                    0xFD -> hwInfo[eventType] = "${r.bits(24, 4)}.${r.bits(28, 4)}"
                    0xFE -> hwInfo[eventType] = "${r.bits(24, 4)}.${r.bits(28, 4)}"
                }
                if (hwInfo.size == 4) {
                    val name = hwInfo[0xFC]
                    listOf(SmartCubeEvent.Hardware(name, hwInfo[0xFE], hwInfo[0xFD], hwInfo[0xFA], name in GYRO_MODELS))
                } else emptyList()
            }
            0xEC -> listOf(gyro(r, 16))
            0xEF -> listOf(SmartCubeEvent.Battery(minOf(r.bits(8 + dataLength * 8, 8), 100)))
            0xEA -> { conn.disconnect(); emptyList() }
            else -> emptyList()
        }
    }

    private companion object {
        val GYRO_MODELS = setOf("GAN12uiM")
    }
}