package coffee.liz.rookery.bt

import coffee.liz.rookery.bt.SmartCubeEvent.Battery
import coffee.liz.rookery.bt.SmartCubeEvent.Facelets
import coffee.liz.rookery.bt.SmartCubeEvent.Gyro
import coffee.liz.rookery.bt.SmartCubeEvent.Hardware
import coffee.liz.rookery.bt.SmartCubeEvent.Move
import kotlinx.coroutines.test.runTest
import kotlin.math.abs
import kotlin.test.Test
import kotlin.test.assertContentEquals
import kotlin.test.assertEquals
import kotlin.test.assertTrue

/** Records commands a driver sends back to the cube. */
class RecordingRawConnection : GanRawConnection {
    val sent = mutableListOf<ByteArray>()
    var disconnects = 0
    override suspend fun sendCommandMessage(message: ByteArray) { sent += message }
    override suspend fun disconnect() { disconnects++ }
}

/** Cube state after F R (from gan-web-bluetooth's docs). */
val F_R_STATE = CubeState(
    cp = listOf(0, 5, 2, 1, 7, 4, 6, 3), co = listOf(1, 2, 0, 2, 1, 1, 0, 2),
    ep = listOf(1, 9, 2, 3, 11, 8, 6, 7, 4, 5, 10, 0), eo = listOf(1, 1, 0, 0, 0, 1, 0, 0, 0, 1, 0, 0),
)
const val F_R_FACELETS = "UUFUUFLLFUUURRRRRRFFRFFDFFDRRBDDBDDBLLDLLDLLDLBBUBBUBB"
val SOLVED_STATE = CubeState(List(8) { it }, List(8) { 0 }, List(12) { it }, List(12) { 0 })

/** Writes the first 7 corners / 11 edges of [s] at the given bit offsets, as the cube does. */
fun GanBitWriter.state(s: CubeState, cpAt: Int, coAt: Int, epAt: Int, eoAt: Int): GanBitWriter {
    for (i in 0 until 7) { put(cpAt + i * 3, 3, s.cp[i]); put(coAt + i * 2, 2, s.co[i]) }
    for (i in 0 until 11) { put(epAt + i * 4, 4, s.ep[i]); put(eoAt + i, 1, s.eo[i]) }
    return this
}

object Packets {
    fun gen2Facelets(serial: Int, s: CubeState) = GanBitWriter(20).put(0, 4, 0x4).put(4, 8, serial).state(s, 12, 33, 47, 91).bytes

    /** [moves] newest first: (face 0..5 URFDLB, direction 0/1, elapsed ms since previous move). */
    fun gen2Moves(serial: Int, moves: List<Triple<Int, Int, Int>>): ByteArray {
        val w = GanBitWriter(20).put(0, 4, 0x2).put(4, 8, serial)
        moves.forEachIndexed { i, (face, dir, elapsed) -> w.put(12 + 5 * i, 4, face).put(16 + 5 * i, 1, dir).put(47 + 16 * i, 16, elapsed) }
        return w.bytes
    }

    // Live-move face codes for Gen3/Gen4, indexed by face URFDLB.
    val LIVE = listOf(2, 32, 8, 1, 16, 4)
    val HISTORY = listOf(1, 5, 3, 0, 4, 2)

    fun gen3Facelets(serial: Int, s: CubeState) =
        GanBitWriter(20).put(0, 8, 0x55).put(8, 8, 0x02).put(16, 8, 0x10).put(24, 16, serial, true).state(s, 40, 61, 77, 121).bytes

    fun gen3Move(serial: Int, face: Int, dir: Int, cubeTs: Long) =
        GanBitWriter(16).put(0, 8, 0x55).put(8, 8, 0x01).put(16, 8, 0x08).put(24, 32, cubeTs, true)
            .put(56, 16, serial, true).put(72, 2, dir).put(74, 6, LIVE[face]).bytes

    fun gen4Facelets(serial: Int, s: CubeState) =
        GanBitWriter(20).put(0, 8, 0xED).put(8, 8, 0x10).put(16, 16, serial, true).state(s, 32, 53, 69, 113).bytes

    fun gen4Move(serial: Int, face: Int, dir: Int, cubeTs: Long) =
        GanBitWriter(20).put(0, 8, 0x01).put(8, 8, 0x08).put(16, 32, cubeTs, true)
            .put(48, 16, serial, true).put(64, 2, dir).put(66, 6, LIVE[face]).bytes

    /** Gen4 move history: [moves] are (face, dir) for serials startSerial, startSerial-1, ... */
    fun gen4History(startSerial: Int, moves: List<Pair<Int, Int>>): ByteArray {
        val w = GanBitWriter(20).put(0, 8, 0xD1).put(8, 8, moves.size / 2 + 1).put(16, 8, startSerial)
        moves.forEachIndexed { i, (face, dir) -> w.put(24 + 4 * i, 3, HISTORY[face]).put(27 + 4 * i, 1, dir) }
        return w.bytes
    }
}

class GanProtocolTest {
    private var now = 1_000L
    private val clock = { now }

    @Test fun bitReaderRoundTripsAllWidths() {
        val w = GanBitWriter(12).put(3, 5, 0b10110).put(8, 16, 0xBEEF).put(27, 16, 0x1234, true).put(43, 32, 0xDEADBEEFL, true)
        val r = GanBitReader(w.bytes)
        assertEquals(0b10110, r.bits(3, 5))
        assertEquals(0xBEEF, r.bits(8, 16))
        assertEquals(0x1234, r.bits(27, 16, true))
        assertEquals(0xDEADBEEFL, r.u32(43, true))
        assertEquals(0, r.bits(200, 8)) // past the end reads as zero
    }

    @Test fun kociembaFacelets() {
        assertEquals(GanFacelets.SOLVED, GanFacelets.toKociemba(SOLVED_STATE))
        assertEquals(F_R_FACELETS, GanFacelets.toKociemba(F_R_STATE))
        val completed = GanFacelets.complete(F_R_STATE.cp.take(7), F_R_STATE.co.take(7), F_R_STATE.ep.take(11), F_R_STATE.eo.take(11))
        assertEquals(F_R_STATE, completed)
    }

    @Test fun gen2Commands() {
        val d = GanGen2ProtocolDriver(clock)
        assertEquals(20, d.createCommandMessage(GanCommand.REQUEST_FACELETS).size)
        assertEquals(0x04, d.createCommandMessage(GanCommand.REQUEST_FACELETS)[0].toInt())
        assertEquals(0x05, d.createCommandMessage(GanCommand.REQUEST_HARDWARE)[0].toInt())
        assertEquals(0x09, d.createCommandMessage(GanCommand.REQUEST_BATTERY)[0].toInt())
        assertContentEquals(
            byteArrayOf(0x0A, 0x05, 0x39, 0x77, 0, 0, 0x01, 0x23, 0x45, 0x67, 0x89.toByte(), 0xAB.toByte(), 0, 0, 0, 0, 0, 0, 0, 0),
            d.createCommandMessage(GanCommand.REQUEST_RESET),
        )
    }

    @Test fun gen2FaceletsMovesBatteryHardwareGyro() = runTest {
        val d = GanGen2ProtocolDriver(clock)
        val c = RecordingRawConnection()
        // Moves before the first facelets event are ignored.
        assertEquals(emptyList(), d.handleStateEvent(c, Packets.gen2Moves(5, listOf(Triple(1, 0, 100)))))

        val f = d.handleStateEvent(c, Packets.gen2Facelets(10, F_R_STATE)).single() as Facelets
        assertEquals(F_R_FACELETS, f.facelets)
        assertEquals(10, f.serial)
        assertEquals(F_R_STATE, f.state)

        now = 2_000
        // Serial 11: one new move R (face 1, CW), 120 ms after the previous.
        val m1 = d.handleStateEvent(c, Packets.gen2Moves(11, listOf(Triple(1, 0, 120)))).single() as Move
        assertEquals(Move("R", 120, 2_000, 11, recovered = false), m1)

        // Serial 13: two moves; the older (U', serial 12) was missed live and comes with the newer (F, serial 13).
        now = 2_500
        val ms = d.handleStateEvent(c, Packets.gen2Moves(13, listOf(Triple(2, 0, 200), Triple(0, 1, 150))))
        assertEquals(listOf(Move("U'", 270, 2_500, 12, true), Move("F", 470, 2_500, 13, false)), ms)

        val battery = GanBitWriter(20).put(0, 4, 0x9).put(8, 8, 87).bytes
        assertEquals(Battery(87), d.handleStateEvent(c, battery).single())
        assertEquals(Battery(100), d.handleStateEvent(c, GanBitWriter(20).put(0, 4, 0x9).put(8, 8, 140).bytes).single())

        val hw = GanBitWriter(20).put(0, 4, 0x5).put(8, 8, 1).put(16, 8, 2).put(24, 8, 3).put(32, 8, 4).put(104, 1, 1)
        "GAN12ui!".forEachIndexed { i, ch -> hw.put(40 + 8 * i, 8, ch.code) }
        assertEquals(Hardware("GAN12ui!", "1.2", "3.4", null, true), d.handleStateEvent(c, hw.bytes).single())

        // Gyro: w = +1, x = -0.5 (sign bit + magnitude), velocity x = -3.
        val g = GanBitWriter(20).put(0, 4, 0x1).put(4, 16, 0x7FFF).put(20, 16, 0x8000 or 0x3FFF).put(68, 4, 0b1011)
        val gyro = d.handleStateEvent(c, g.bytes).single() as Gyro
        assertEquals(1.0, gyro.quaternion.w)
        assertTrue(abs(gyro.quaternion.x + 0.5) < 1e-3)
        assertEquals(-3, gyro.velocity!!.x)

        assertEquals(emptyList(), d.handleStateEvent(c, GanBitWriter(20).put(0, 4, 0xD).bytes))
        assertEquals(1, c.disconnects)
    }

    @Test fun gen3Commands() {
        val d = GanGen3ProtocolDriver(clock)
        assertContentEquals(byteArrayOf(0x68, 0x01) + ByteArray(14), d.createCommandMessage(GanCommand.REQUEST_FACELETS))
        assertContentEquals(byteArrayOf(0x68, 0x04) + ByteArray(14), d.createCommandMessage(GanCommand.REQUEST_HARDWARE))
        assertContentEquals(byteArrayOf(0x68, 0x07) + ByteArray(14), d.createCommandMessage(GanCommand.REQUEST_BATTERY))
        assertEquals(0x05, d.createCommandMessage(GanCommand.REQUEST_RESET)[2].toInt())
    }

    @Test fun gen3FaceletsMovesBatteryHardware() = runTest {
        val d = GanGen3ProtocolDriver(clock)
        val c = RecordingRawConnection()
        assertEquals(emptyList(), d.handleStateEvent(c, Packets.gen3Move(1, 0, 0, 10)))
        assertEquals(Facelets(GanFacelets.SOLVED, 7, SOLVED_STATE), d.handleStateEvent(c, Packets.gen3Facelets(7, SOLVED_STATE)).single())
        assertEquals(Move("L'", 123_456, 1_000, 8), d.handleStateEvent(c, Packets.gen3Move(8, 4, 1, 123_456)).single())
        assertEquals(Move("B", 0xFFFFFFF0L, 1_000, 9), d.handleStateEvent(c, Packets.gen3Move(9, 5, 0, 0xFFFFFFF0L)).single())
        // Wrong magic is ignored.
        assertEquals(emptyList(), d.handleStateEvent(c, Packets.gen3Move(10, 0, 0, 1).also { it[0] = 0x54 }))

        val battery = GanBitWriter(16).put(0, 8, 0x55).put(8, 8, 0x10).put(16, 8, 1).put(24, 8, 64).bytes
        assertEquals(Battery(64), d.handleStateEvent(c, battery).single())

        val hw = GanBitWriter(16).put(0, 8, 0x55).put(8, 8, 0x07).put(16, 8, 9).put(72, 4, 2).put(76, 4, 5).put(80, 4, 1).put(84, 4, 0)
        "GANic".forEachIndexed { i, ch -> hw.put(32 + 8 * i, 8, ch.code) }
        assertEquals(Hardware("GANic", "1.0", "2.5", null, false), d.handleStateEvent(c, hw.bytes).single())
    }

    @Test fun gen4Commands() {
        val d = GanGen4ProtocolDriver(clock)
        assertContentEquals(byteArrayOf(0xDD.toByte(), 0x04, 0x00, 0xED.toByte(), 0, 0) + ByteArray(14), d.createCommandMessage(GanCommand.REQUEST_FACELETS))
        assertContentEquals(byteArrayOf(0xDD.toByte(), 0x04, 0x00, 0xEF.toByte(), 0, 0) + ByteArray(14), d.createCommandMessage(GanCommand.REQUEST_BATTERY))
        assertContentEquals(byteArrayOf(0xDF.toByte(), 0x03) + ByteArray(18), d.createCommandMessage(GanCommand.REQUEST_HARDWARE))
        assertEquals(0xD2.toByte(), d.createCommandMessage(GanCommand.REQUEST_RESET)[0])
    }

    @Test fun gen4MovesRecoverMissedMovesFromHistory() = runTest {
        val d = GanGen4ProtocolDriver(clock)
        val c = RecordingRawConnection()
        assertEquals(F_R_FACELETS, (d.handleStateEvent(c, Packets.gen4Facelets(10, F_R_STATE)).single() as Facelets).facelets)
        assertEquals(listOf(Move("D", 5000, 1_000, 11)), d.handleStateEvent(c, Packets.gen4Move(11, 3, 0, 5000)))

        // Serial 14 arrives but 12 and 13 were lost: it is held back and history is requested.
        assertEquals(emptyList(), d.handleStateEvent(c, Packets.gen4Move(14, 1, 1, 5600)))
        assertContentEquals(byteArrayOf(0xD1.toByte(), 0x04, 13, 0, 4, 0) + ByteArray(14), c.sent.single())

        // History: serials 13 (F'), 12 (U), 11 (D), 10 (whatever) newest first.
        now = 1_050
        val recovered = d.handleStateEvent(c, Packets.gen4History(13, listOf(2 to 1, 0 to 0, 3 to 0, 5 to 0)))
        assertEquals(
            listOf(Move("U", null, 1_050, 12, true), Move("F'", null, 1_050, 13, true), Move("R'", 5600, 1_000, 14)),
            recovered,
        )
    }

    @Test fun gen4RetriesHistoryThenGivesUpInsteadOfHanging() = runTest {
        // The failure seen live on a GANicE3: a move after a lost packet waited until the next turn.
        val d = GanGen4ProtocolDriver(clock)
        val c = RecordingRawConnection()
        d.handleStateEvent(c, Packets.gen4Facelets(10, SOLVED_STATE))
        d.handleStateEvent(c, Packets.gen4Move(11, 4, 1, 100))
        assertEquals(emptyList(), d.handleStateEvent(c, Packets.gen4Move(13, 5, 0, 155))) // 12 lost
        assertEquals(1, c.sent.size)

        now += 100; assertEquals(emptyList(), d.tick(c)); assertEquals(1, c.sent.size) // too soon to retry
        now += 250; assertEquals(emptyList(), d.tick(c)); assertEquals(2, c.sent.size) // retried
        assertEquals(0xD1.toByte(), c.sent.last()[0])

        // History never comes: after 1.5 s the held move is released, flagged, and facelets requested.
        now += 1_200
        val released = d.tick(c)
        assertEquals(listOf(Move("B", 155, 1_000, 13, afterGap = true)), released)
        assertContentEquals(byteArrayOf(0xDD.toByte(), 0x04, 0x00, 0xED.toByte(), 0x00, 0x00) + ByteArray(14), c.sent.last())
        assertEquals(0, c.disconnects)

        // Live moves flow again immediately afterwards.
        assertEquals(listOf("U"), d.handleStateEvent(c, Packets.gen4Move(14, 0, 0, 900)).map { (it as Move).move })
        assertEquals(emptyList(), d.tick(c))
    }

    @Test fun gen4BigBacklogIsReleasedNotDisconnected() = runTest {
        val d = GanGen4ProtocolDriver(clock)
        val c = RecordingRawConnection()
        d.handleStateEvent(c, Packets.gen4Facelets(10, SOLVED_STATE))
        d.handleStateEvent(c, Packets.gen4Move(11, 0, 0, 1))
        val out = (13..30).flatMap { d.handleStateEvent(c, Packets.gen4Move(it, 0, 0, it.toLong())) }
        assertEquals((13..30).toList(), out.map { (it as Move).serial }) // 13-29 released, then 30 flows live
        assertTrue((out.first() as Move).afterGap)
        assertEquals(0, c.disconnects)
    }

    @Test fun gen4FaceletsTriggerHistoryWhenIdle() = runTest {
        val d = GanGen4ProtocolDriver(clock)
        val c = RecordingRawConnection()
        d.handleStateEvent(c, Packets.gen4Facelets(20, SOLVED_STATE))
        d.handleStateEvent(c, Packets.gen4Move(21, 0, 0, 100))
        // Periodic facelets says serial 23 more than 500 ms after the last move: 22 and 23 were missed.
        now += 600
        d.handleStateEvent(c, Packets.gen4Facelets(23, SOLVED_STATE))
        assertContentEquals(byteArrayOf(0xD1.toByte(), 0x04, 23, 0, 4, 0) + ByteArray(14), c.sent.single())
        val got = d.handleStateEvent(c, Packets.gen4History(23, listOf(1 to 0, 2 to 0, 0 to 0, 0 to 0)))
        assertEquals(listOf("F", "R"), got.map { (it as Move).move })
        assertEquals(listOf(22, 23), got.map { (it as Move).serial })
    }

    @Test fun gen4BatteryHardwareGyro() = runTest {
        val d = GanGen4ProtocolDriver(clock)
        val c = RecordingRawConnection()
        // Battery level sits right after the payload: byte 1 + dataLength.
        val battery = GanBitWriter(20).put(0, 8, 0xEF).put(8, 8, 3).put(32, 8, 55).bytes
        assertEquals(Battery(55), d.handleStateEvent(c, battery).single())

        d.createCommandMessage(GanCommand.REQUEST_HARDWARE)
        val date = GanBitWriter(20).put(0, 8, 0xFA).put(8, 8, 5).put(24, 16, 2023, true).put(40, 8, 7).put(48, 8, 9).bytes
        val nameW = GanBitWriter(20).put(0, 8, 0xFC).put(8, 8, 9)
        "GAN12uiM".forEachIndexed { i, ch -> nameW.put(24 + 8 * i, 8, ch.code) }
        val sw = GanBitWriter(20).put(0, 8, 0xFD).put(8, 8, 2).put(24, 4, 1).put(28, 4, 6).bytes
        val hw = GanBitWriter(20).put(0, 8, 0xFE).put(8, 8, 2).put(24, 4, 1).put(28, 4, 1).bytes
        assertEquals(emptyList(), d.handleStateEvent(c, date))
        assertEquals(emptyList(), d.handleStateEvent(c, nameW.bytes))
        assertEquals(emptyList(), d.handleStateEvent(c, sw))
        assertEquals(Hardware("GAN12uiM", "1.1", "1.6", "2023-07-09", true), d.handleStateEvent(c, hw).single())

        val g = GanBitWriter(20).put(0, 8, 0xEC).put(8, 8, 0x0A).put(16, 16, 0x7FFF).put(48, 16, 0x8000 or 0x7FFF).put(88, 4, 0b0010).bytes
        val gyro = d.handleStateEvent(c, g).single() as Gyro
        assertEquals(Quaternion(0.0, -1.0, 0.0, 1.0), gyro.quaternion)
        assertEquals(AngularVelocity(0, 0, 2), gyro.velocity)
    }

    @Test fun timestampLinearFit() {
        // Cube clock runs 10% slow relative to the host.
        val fit = GanTimestampFit.linearFit(listOf(TimedMove("R", 1000, 0), TimedMove("U", 1110, 100), TimedMove("R'", 1220, 200)))
        assertEquals(listOf(0L, 110L, 220L), fit.map { it.cubeTimestampMs })
        assertEquals(-9.091, GanTimestampFit.skewPercent(listOf(TimedMove("R", 1000, 0), TimedMove("U", 1110, 100), TimedMove("R'", 1220, 200))))

        // A recovered move without timestamps is interpolated from its neighbours (next - 50 ms).
        val gap = GanTimestampFit.linearFit(listOf(TimedMove("R", 1000, 0), TimedMove("U", null, null), TimedMove("R'", 1220, 200)))
        assertEquals(listOf(0L, 165L, 220L), gap.map { it.cubeTimestampMs })
        val lastMissing = GanTimestampFit.linearFit(listOf(TimedMove("R", 1000, 0), TimedMove("U", 1100, 100), TimedMove("R'", null, null)))
        assertEquals(listOf(0L, 100L, 150L), lastMissing.map { it.cubeTimestampMs })
        assertEquals(emptyList(), GanTimestampFit.linearFit(emptyList()))
    }

    @Test fun macAndScanHelpers() {
        assertContentEquals(byteArrayOf(0x78, 0x56, 0x34, 0x12, 0xF0.toByte(), 0x6F), GanDefinitions.saltFromMac("6F:F0:12:34:56:78"))
        assertContentEquals(GanDefinitions.saltFromMac("6f-f0-12-34-56-78"), GanDefinitions.saltFromMac("6F:F0:12:34:56:78"))
        // Manufacturer data (company 0x1201): 3 bytes of something, then the MAC reversed.
        val mfr = mapOf(0x1201 to byteArrayOf(1, 2, 3, 0x78, 0x56, 0x34, 0x12, 0xF0.toByte(), 0x6F, 9, 9))
        assertEquals("6F:F0:12:34:56:78", GanDefinitions.macFromManufacturerData(mfr))
        assertTrue(GanCubeConnector.isGanCube(BleAdvertisement("6F:F0:12:34:56:78", null, -60, mfr)))
        assertTrue(GanCubeConnector.isGanCube(BleAdvertisement("00:11:22:33:44:55", "GAN12uiM_1A2B", -60)))
        assertTrue(GanCubeConnector.isGanCube(BleAdvertisement("00:11:22:33:44:55", null, -60, serviceUuids = listOf(GanDefinitions.GEN3_SERVICE))))
        assertTrue(!GanCubeConnector.isGanCube(BleAdvertisement("00:11:22:33:44:55", "Headphones", -60, mfr)))
        assertEquals(GanGeneration.GEN4, GanGeneration.fromServiceUuids(listOf("0000180a-0000-1000-8000-00805f9b34fb", GanDefinitions.GEN4_SERVICE.uppercase())))
    }
}