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

/** GATT UUIDs and key material for GAN smart cubes. */
object GanDefinitions {
    const val GEN2_SERVICE = "6e400001-b5a3-f393-e0a9-e50e24dc4179"
    const val GEN2_COMMAND = "28be4a4a-cd67-11e9-a32f-2a2ae2dbcce4"
    const val GEN2_STATE = "28be4cb6-cd67-11e9-a32f-2a2ae2dbcce4"

    const val GEN3_SERVICE = "8653000a-43e6-47b7-9cb0-5fc21d4ae340"
    const val GEN3_COMMAND = "8653000c-43e6-47b7-9cb0-5fc21d4ae340"
    const val GEN3_STATE = "8653000b-43e6-47b7-9cb0-5fc21d4ae340"

    const val GEN4_SERVICE = "00000010-0000-fff7-fff6-fff5fff4fff0"
    const val GEN4_COMMAND = "0000fff5-0000-1000-8000-00805f9b34fb"
    const val GEN4_STATE = "0000fff6-0000-1000-8000-00805f9b34fb"

    /** Advertised name prefixes of GAN-protocol cubes (MoYu AI 2023 = "AiCube" uses Gen2 with its own key). */
    val NAME_PREFIXES = listOf("GAN", "MG", "AiCube")

    /** Company identifier codes GAN cubes use: 0x0001, 0x0101, ... 0xFF01. */
    fun isGanCic(cic: Int): Boolean = (cic and 0xFF) == 0x01

    class KeyIv(val key: ByteArray, val iv: ByteArray)

    private fun bytes(vararg v: Int) = ByteArray(v.size) { v[it].toByte() }

    /** Key used by GAN Gen2, Gen3 and Gen4 cubes. */
    val GAN_KEY = KeyIv(
        bytes(0x01, 0x02, 0x42, 0x28, 0x31, 0x91, 0x16, 0x07, 0x20, 0x05, 0x18, 0x54, 0x42, 0x11, 0x12, 0x53),
        bytes(0x11, 0x03, 0x32, 0x28, 0x21, 0x01, 0x76, 0x27, 0x20, 0x95, 0x78, 0x14, 0x32, 0x12, 0x02, 0x43),
    )

    /** Key used by MoYu AI 2023 (advertises as "AiCube", speaks GAN Gen2). */
    val MOYU_AI_KEY = KeyIv(
        bytes(0x05, 0x12, 0x02, 0x45, 0x02, 0x01, 0x29, 0x56, 0x12, 0x78, 0x12, 0x76, 0x81, 0x01, 0x08, 0x03),
        bytes(0x01, 0x44, 0x28, 0x06, 0x86, 0x21, 0x22, 0x28, 0x51, 0x05, 0x08, 0x31, 0x82, 0x02, 0x21, 0x06),
    )

    /** Encryption salt: the 6 MAC bytes in reverse order. Accepts "AA:BB:...", "AA-BB-..." or spaces. */
    fun saltFromMac(mac: String): ByteArray {
        val parts = mac.trim().split(Regex("[:\\-\\s]+"))
        require(parts.size == 6) { "MAC must have 6 bytes: $mac" }
        return ByteArray(6) { parts[5 - it].toInt(16).toByte() }
    }

    /**
     * The cube's MAC from GAN manufacturer-specific data: the last 6 of the first 9 payload bytes, reversed.
     * Returns null if no GAN company id is present.
     */
    fun macFromManufacturerData(data: Map<Int, ByteArray>): String? {
        val payload = data.entries.firstOrNull { isGanCic(it.key) }?.value ?: return null
        val view = payload.copyOfRange(0, minOf(9, payload.size))
        if (view.size < 6) return null
        return (1..6).joinToString(":") { i -> (view[view.size - i].toInt() and 0xFF).toString(16).uppercase().padStart(2, '0') }
    }
}

/**
 * GAN's message encryption (identical for Gen2/3/4): AES-128-CBC with a fresh IV on each 16-byte chunk,
 * applied to the chunk at the start of the message and then to the (possibly overlapping) chunk at its end.
 * Key and IV are salted with the reversed MAC: `k[i] = (k[i] + salt[i]) % 0xFF` for i < 6
 * (sic: modulo 255, which is what the cube firmware does).
 */
class GanEncrypter(private val aes: Aes128, baseKey: ByteArray, baseIv: ByteArray, salt: ByteArray) {
    val key: ByteArray = baseKey.copyOf()
    val iv: ByteArray = baseIv.copyOf()

    init {
        require(baseKey.size == 16) { "Key must be 16 bytes (128-bit) long" }
        require(baseIv.size == 16) { "Iv must be 16 bytes (128-bit) long" }
        require(salt.size == 6) { "Salt must be 6 bytes (48-bit) long" }
        for (i in 0 until 6) {
            key[i] = (((baseKey[i].toInt() and 0xFF) + (salt[i].toInt() and 0xFF)) % 0xFF).toByte()
            iv[i] = (((baseIv[i].toInt() and 0xFF) + (salt[i].toInt() and 0xFF)) % 0xFF).toByte()
        }
    }

    constructor(aes: Aes128, keyIv: GanDefinitions.KeyIv, mac: String) :
        this(aes, keyIv.key, keyIv.iv, GanDefinitions.saltFromMac(mac))

    private fun encryptChunk(buf: ByteArray, off: Int) {
        val block = ByteArray(16) { (buf[off + it].toInt() xor iv[it].toInt()).toByte() }
        aes.encryptBlock(key, block).copyInto(buf, off)
    }

    private fun decryptChunk(buf: ByteArray, off: Int) {
        val plain = aes.decryptBlock(key, buf.copyOfRange(off, off + 16))
        for (i in 0 until 16) buf[off + i] = (plain[i].toInt() xor iv[i].toInt()).toByte()
    }

    fun encrypt(data: ByteArray): ByteArray {
        require(data.size >= 16) { "Data must be at least 16 bytes long" }
        val res = data.copyOf()
        encryptChunk(res, 0)
        if (res.size > 16) encryptChunk(res, res.size - 16)
        return res
    }

    fun decrypt(data: ByteArray): ByteArray {
        require(data.size >= 16) { "Data must be at least 16 bytes long" }
        val res = data.copyOf()
        if (res.size > 16) decryptChunk(res, res.size - 16)
        decryptChunk(res, 0)
        return res
    }
}