package no.zyon.shade.serialization

import no.zyon.shade.types.PreKeyMessage
import no.zyon.shade.types.RatchetMessage
import no.zyon.shade.types.ShadeEnvelope
import java.nio.ByteBuffer

/**
 * Compact binary wire format. MUST match @shade/proto/wire.ts byte-for-byte.
 *
 * Format: [version:1][type:1][payload...]
 * Types: 0x01 = PreKeyMessage, 0x02 = RatchetMessage
 * Integers: big-endian
 * Byte arrays: 4-byte (u32) length prefix + data (since wire VERSION 0x02).
 *
 * VERSION 0x01 used a 2-byte length prefix and was capped at 64 KiB
 * payloads — incompatible with inline file ops up to 256 KiB.
 */
object WireFormat {
    private const val VERSION: Byte = 0x02
    private const val TYPE_PREKEY: Byte = 0x01
    private const val TYPE_RATCHET: Byte = 0x02
    private const val PREKEY_NONE: Long = 0xFFFFFFFFL

    // ─── Encode ──────────────────────────────────────────────

    fun encodeEnvelope(envelope: ShadeEnvelope): ByteArray {
        return when (envelope.type) {
            ShadeEnvelope.EnvelopeType.PREKEY ->
                encodePreKeyMessage(envelope.content as PreKeyMessage)
            ShadeEnvelope.EnvelopeType.RATCHET ->
                encodeRatchetMessage(envelope.content as RatchetMessage)
        }
    }

    fun encodePreKeyMessage(msg: PreKeyMessage): ByteArray {
        val ratchetBytes = encodeRatchetInner(msg.message)
        val parts = mutableListOf<ByteArray>()
        parts.add(byteArrayOf(VERSION, TYPE_PREKEY))
        parts.add(uint32(msg.registrationId.toLong()))
        parts.add(uint32(msg.preKeyId?.toLong() ?: PREKEY_NONE))
        parts.add(uint32(msg.signedPreKeyId.toLong()))
        parts.add(lpBytes(msg.ephemeralKey))
        parts.add(lpBytes(msg.identityDHKey))
        parts.add(lpBytes(ratchetBytes))
        return concat(parts)
    }

    fun encodeRatchetMessage(msg: RatchetMessage): ByteArray {
        val parts = mutableListOf<ByteArray>()
        parts.add(byteArrayOf(VERSION, TYPE_RATCHET))
        parts.add(encodeRatchetInner(msg))
        return concat(parts)
    }

    private fun encodeRatchetInner(msg: RatchetMessage): ByteArray {
        val parts = mutableListOf<ByteArray>()
        parts.add(lpBytes(msg.dhPublicKey))
        parts.add(uint32(msg.previousCounter.toLong()))
        parts.add(uint32(msg.counter.toLong()))
        parts.add(lpBytes(msg.ciphertext))
        parts.add(lpBytes(msg.nonce))
        return concat(parts)
    }

    // ─── Decode ──────────────────────────────────────────────

    fun decodeEnvelope(data: ByteArray): ShadeEnvelope {
        if (data.size < 2) throw IllegalArgumentException("Too short")
        val version = data[0]
        if (version != VERSION) throw IllegalArgumentException("Unknown version: $version")
        val type = data[1]
        val payload = data.copyOfRange(2, data.size)

        return when (type) {
            TYPE_PREKEY -> ShadeEnvelope(
                type = ShadeEnvelope.EnvelopeType.PREKEY,
                content = decodePreKeyMessageInner(payload),
                timestamp = 0,
                senderAddress = "",
            )
            TYPE_RATCHET -> {
                val (msg, _) = decodeRatchetInner(payload, 0)
                ShadeEnvelope(
                    type = ShadeEnvelope.EnvelopeType.RATCHET,
                    content = msg,
                    timestamp = 0,
                    senderAddress = "",
                )
            }
            else -> throw IllegalArgumentException("Unknown type: $type")
        }
    }

    private fun decodePreKeyMessageInner(data: ByteArray): PreKeyMessage {
        var offset = 0
        val registrationId = readUint32(data, offset).toInt(); offset += 4
        val preKeyIdRaw = readUint32(data, offset); offset += 4
        val preKeyId = if (preKeyIdRaw == PREKEY_NONE) null else preKeyIdRaw.toInt()
        val signedPreKeyId = readUint32(data, offset).toInt(); offset += 4

        val ephemeral = readLP(data, offset); offset = ephemeral.second
        val identityDH = readLP(data, offset); offset = identityDH.second
        val ratchetData = readLP(data, offset); offset = ratchetData.second

        val (ratchet, _) = decodeRatchetInner(ratchetData.first, 0)

        return PreKeyMessage(
            registrationId = registrationId,
            preKeyId = preKeyId,
            signedPreKeyId = signedPreKeyId,
            ephemeralKey = ephemeral.first,
            identityDHKey = identityDH.first,
            message = ratchet,
        )
    }

    private fun decodeRatchetInner(data: ByteArray, startOffset: Int): Pair<RatchetMessage, Int> {
        var offset = startOffset
        val dhPub = readLP(data, offset); offset = dhPub.second
        val prevCounter = readUint32(data, offset).toInt(); offset += 4
        val counter = readUint32(data, offset).toInt(); offset += 4
        val ciphertext = readLP(data, offset); offset = ciphertext.second
        val nonce = readLP(data, offset); offset = nonce.second

        return RatchetMessage(
            dhPublicKey = dhPub.first,
            previousCounter = prevCounter,
            counter = counter,
            ciphertext = ciphertext.first,
            nonce = nonce.first,
        ) to offset
    }

    // ─── Helpers ─────────────────────────────────────────────

    private fun uint32(n: Long): ByteArray {
        val buf = ByteBuffer.allocate(4)
        buf.putInt(n.toInt())
        return buf.array()
    }

    private fun lpBytes(data: ByteArray): ByteArray {
        val len = ByteBuffer.allocate(4)
        len.putInt(data.size)
        return concat(listOf(len.array(), data))
    }

    private fun readUint32(data: ByteArray, offset: Int): Long {
        return ((data[offset].toLong() and 0xff) shl 24) or
            ((data[offset + 1].toLong() and 0xff) shl 16) or
            ((data[offset + 2].toLong() and 0xff) shl 8) or
            (data[offset + 3].toLong() and 0xff)
    }

    private fun readLP(data: ByteArray, offset: Int): Pair<ByteArray, Int> {
        val len = readUint32(data, offset).toInt()
        val value = data.copyOfRange(offset + 4, offset + 4 + len)
        return value to (offset + 4 + len)
    }

    private fun concat(parts: List<ByteArray>): ByteArray {
        val total = parts.sumOf { it.size }
        val result = ByteArray(total)
        var offset = 0
        for (p in parts) {
            p.copyInto(result, offset)
            offset += p.size
        }
        return result
    }
}
