package no.zyon.shade.serialization

import no.zyon.shade.streams.StreamConstants
import java.nio.ByteBuffer

/**
 * Wire-decoded stream-chunk envelope (type 0x11).
 *
 * Mirror @shade/proto/wire.ts `StreamChunkWire`. The nonce is deterministic
 * (derived from `(laneId, seq)` on both sides) but is also serialized over
 * the wire for self-description and validated by the receiver.
 *
 * `seq` is unsigned-u64 on the wire; on the JVM we keep it as Long. The
 * encode/decode helpers operate on the raw 8-byte big-endian representation,
 * so values past Long.MAX_VALUE roundtrip via `Long.toULong()`.
 */
data class StreamChunkWire(
    val streamId: ByteArray,
    val laneId: Long,
    val seq: Long,
    val isLast: Boolean,
    val nonce: ByteArray,
    val aad: ByteArray,
    val ciphertext: ByteArray,
) {
    override fun equals(other: Any?): Boolean {
        if (this === other) return true
        if (other !is StreamChunkWire) return false
        return streamId.contentEquals(other.streamId) &&
            laneId == other.laneId &&
            seq == other.seq &&
            isLast == other.isLast &&
            nonce.contentEquals(other.nonce) &&
            aad.contentEquals(other.aad) &&
            ciphertext.contentEquals(other.ciphertext)
    }

    override fun hashCode(): Int {
        var result = streamId.contentHashCode()
        result = 31 * result + laneId.hashCode()
        result = 31 * result + seq.hashCode()
        result = 31 * result + isLast.hashCode()
        result = 31 * result + nonce.contentHashCode()
        result = 31 * result + aad.contentHashCode()
        result = 31 * result + ciphertext.contentHashCode()
        return result
    }
}

/** Stream-chunk wire codec. Mirror @shade/proto/wire.ts `encodeStreamChunk`/`decodeStreamChunk`. */
object StreamChunkWireFormat {
    private const val VERSION: Byte = 0x02
    const val TYPE_STREAM_CHUNK: Byte = 0x11

    fun encodeStreamChunk(c: StreamChunkWire): ByteArray {
        require(c.streamId.size == StreamConstants.STREAM_ID_BYTES) {
            "streamId must be ${StreamConstants.STREAM_ID_BYTES} bytes"
        }
        require(c.nonce.size == StreamConstants.STREAM_NONCE_BYTES) {
            "nonce must be ${StreamConstants.STREAM_NONCE_BYTES} bytes"
        }
        require(c.laneId in 0L..0xFFFFFFFFL) { "laneId out of u32 range: ${c.laneId}" }
        // c.seq is unsigned-u64; negative signed longs encode as the high half
        // of the u64 range. ByteBuffer.putLong writes the raw 8-byte pattern.

        val headerSize =
            1 + 1 +
                StreamConstants.STREAM_ID_BYTES +
                4 + 8 + 1 +
                StreamConstants.STREAM_NONCE_BYTES +
                4 + c.aad.size +
                4
        val out = ByteArray(headerSize + c.ciphertext.size)
        val buf = ByteBuffer.wrap(out)

        buf.put(VERSION)
        buf.put(TYPE_STREAM_CHUNK)
        buf.put(c.streamId)
        buf.putInt(c.laneId.toInt())
        buf.putLong(c.seq)
        buf.put(if (c.isLast) 0x01.toByte() else 0x00.toByte())
        buf.put(c.nonce)
        buf.putInt(c.aad.size)
        buf.put(c.aad)
        buf.putInt(c.ciphertext.size)
        buf.put(c.ciphertext)

        return out
    }

    fun decodeStreamChunk(data: ByteArray): StreamChunkWire {
        val minHeaderSize = 2 +
            StreamConstants.STREAM_ID_BYTES +
            4 + 8 + 1 +
            StreamConstants.STREAM_NONCE_BYTES +
            4 + 4
        require(data.size >= minHeaderSize) {
            "stream-chunk too short: ${data.size} < $minHeaderSize"
        }
        require(data[0] == VERSION) { "Unknown version: ${data[0]}" }
        require(data[1] == TYPE_STREAM_CHUNK) { "Not a stream-chunk: type=${data[1]}" }

        val buf = ByteBuffer.wrap(data)
        buf.position(2)

        val streamId = ByteArray(StreamConstants.STREAM_ID_BYTES)
        buf.get(streamId)

        val laneId = buf.int.toLong() and 0xFFFFFFFFL
        val seq = buf.long
        val isLast = buf.get() == 0x01.toByte()

        val nonce = ByteArray(StreamConstants.STREAM_NONCE_BYTES)
        buf.get(nonce)

        val aadLen = buf.int
        require(buf.position() + aadLen + 4 <= data.size) {
            "stream-chunk truncated in aad/ctLen"
        }
        val aad = ByteArray(aadLen)
        buf.get(aad)

        val ctLen = buf.int
        require(buf.position() + ctLen == data.size) {
            "stream-chunk length mismatch: declared ${buf.position() + ctLen}, actual ${data.size}"
        }
        val ciphertext = ByteArray(ctLen)
        buf.get(ciphertext)

        return StreamChunkWire(streamId, laneId, seq, isLast, nonce, aad, ciphertext)
    }

    /** Inspect the type tag without full parsing. Mirror @shade/proto/wire.ts. */
    enum class EnvelopeKind { PREKEY, RATCHET, STREAM_CHUNK, UNKNOWN }

    fun inspectEnvelopeType(data: ByteArray): EnvelopeKind {
        if (data.size < 2 || data[0] != VERSION) return EnvelopeKind.UNKNOWN
        return when (data[1]) {
            0x01.toByte() -> EnvelopeKind.PREKEY
            0x02.toByte() -> EnvelopeKind.RATCHET
            TYPE_STREAM_CHUNK -> EnvelopeKind.STREAM_CHUNK
            else -> EnvelopeKind.UNKNOWN
        }
    }
}
