package com.rosetta.messenger.network

/**
 * Binary stream for protocol packets.
 *
 * Parity with desktop/server:
 * - signed: Int8/16/32/64 (two's complement)
 * - unsigned: UInt8/16/32/64
 * - String: length(UInt32) + chars(UInt16)
 * - byte[]: length(UInt32) + raw bytes
 */
class Stream(initial: ByteArray = byteArrayOf()) {
    private var stream: ByteArray = initial.copyOf()
    private var readPointer: Int = 0 // bits
    private var writePointer: Int = stream.size shl 3 // bits

    fun getStream(): ByteArray = stream.copyOf(length())

    fun setStream(value: ByteArray) {
        stream = value.copyOf()
        readPointer = 0
        writePointer = stream.size shl 3
    }

    fun getBuffer(): ByteArray = getStream()

    fun isEmpty(): Boolean = writePointer == 0

    fun length(): Int = (writePointer + 7) shr 3

    fun getReadPointerBits(): Int = readPointer

    fun setReadPointerBits(bits: Int) {
        readPointer = bits.coerceIn(0, writePointer)
    }

    fun getTotalBits(): Int = writePointer

    fun getRemainingBits(): Int = (writePointer - readPointer).coerceAtLeast(0)

    fun hasRemainingBits(): Boolean = readPointer < writePointer

    fun writeBit(value: Int) {
        writeBits((value and 1).toLong(), 1)
    }

    fun readBit(): Int = readBits(1).toInt()

    fun writeBoolean(value: Boolean) {
        writeBit(if (value) 1 else 0)
    }

    fun readBoolean(): Boolean = readBit() == 1

    fun writeByte(value: Byte) {
        writeUInt8(value.toInt() and 0xFF)
    }

    fun readByte(): Byte = readUInt8().toByte()

    fun writeUInt8(value: Int) {
        val v = value and 0xFF

        // Fast path when byte-aligned.
        if ((writePointer and 7) == 0) {
            reserveBits(8)
            stream[writePointer shr 3] = v.toByte()
            writePointer += 8
            return
        }

        writeBits(v.toLong(), 8)
    }

    fun readUInt8(): Int {
        if (remainingBits() < 8L) {
            throw IllegalStateException("Not enough bits to read UInt8")
        }

        // Fast path when byte-aligned.
        if ((readPointer and 7) == 0) {
            val value = stream[readPointer shr 3].toInt() and 0xFF
            readPointer += 8
            return value
        }

        return readBits(8).toInt() and 0xFF
    }

    fun writeInt8(value: Int) {
        writeUInt8(value)
    }

    fun readInt8(): Int = readUInt8().toByte().toInt()

    fun writeUInt16(value: Int) {
        val v = value and 0xFFFF
        writeUInt8((v ushr 8) and 0xFF)
        writeUInt8(v and 0xFF)
    }

    fun readUInt16(): Int {
        val hi = readUInt8()
        val lo = readUInt8()
        return (hi shl 8) or lo
    }

    fun writeInt16(value: Int) {
        writeUInt16(value)
    }

    fun readInt16(): Int = readUInt16().toShort().toInt()

    fun writeUInt32(value: Long) {
        if (value < 0L || value > 0xFFFF_FFFFL) {
            throw IllegalArgumentException("UInt32 out of range: $value")
        }

        writeUInt8(((value ushr 24) and 0xFF).toInt())
        writeUInt8(((value ushr 16) and 0xFF).toInt())
        writeUInt8(((value ushr 8) and 0xFF).toInt())
        writeUInt8((value and 0xFF).toInt())
    }

    fun readUInt32(): Long {
        val b1 = readUInt8().toLong() and 0xFFL
        val b2 = readUInt8().toLong() and 0xFFL
        val b3 = readUInt8().toLong() and 0xFFL
        val b4 = readUInt8().toLong() and 0xFFL
        return (b1 shl 24) or (b2 shl 16) or (b3 shl 8) or b4
    }

    fun writeInt32(value: Int) {
        writeUInt32(value.toLong() and 0xFFFF_FFFFL)
    }

    fun readInt32(): Int = readUInt32().toInt()

    /** Writes raw 64-bit pattern (UInt64 bit-pattern in Long). */
    fun writeUInt64(value: Long) {
        writeUInt8(((value ushr 56) and 0xFF).toInt())
        writeUInt8(((value ushr 48) and 0xFF).toInt())
        writeUInt8(((value ushr 40) and 0xFF).toInt())
        writeUInt8(((value ushr 32) and 0xFF).toInt())
        writeUInt8(((value ushr 24) and 0xFF).toInt())
        writeUInt8(((value ushr 16) and 0xFF).toInt())
        writeUInt8(((value ushr 8) and 0xFF).toInt())
        writeUInt8((value and 0xFF).toInt())
    }

    /** Reads raw 64-bit pattern (UInt64 bit-pattern in Long). */
    fun readUInt64(): Long {
        val high = readUInt32() and 0xFFFF_FFFFL
        val low = readUInt32() and 0xFFFF_FFFFL
        return (high shl 32) or low
    }

    fun writeInt64(value: Long) {
        writeUInt64(value)
    }

    fun readInt64(): Long = readUInt64()

    fun writeFloat32(value: Float) {
        writeInt32(java.lang.Float.floatToIntBits(value))
    }

    fun readFloat32(): Float = java.lang.Float.intBitsToFloat(readInt32())

    /** String: length(UInt32) + chars(UInt16). */
    fun writeString(value: String) {
        writeUInt32(value.length.toLong())

        if (value.isEmpty()) return

        reserveBits(value.length.toLong() * 16L)
        for (i in value.indices) {
            writeUInt16(value[i].code and 0xFFFF)
        }
    }

    fun readString(): String {
        val lenLong = readUInt32()
        if (lenLong > Int.MAX_VALUE.toLong()) {
            throw IllegalStateException("String length too large: $lenLong")
        }

        val length = lenLong.toInt()
        val requiredBits = length.toLong() * 16L
        if (requiredBits > remainingBits()) {
            throw IllegalStateException("Not enough bits to read string")
        }

        val sb = StringBuilder(length)
        repeat(length) {
            sb.append(readUInt16().toChar())
        }
        return sb.toString()
    }

    /** byte[]: length(UInt32) + payload. */
    fun writeBytes(value: ByteArray) {
        writeUInt32(value.size.toLong())

        if (value.isEmpty()) return

        reserveBits(value.size.toLong() * 8L)

        // Fast path when byte-aligned.
        if ((writePointer and 7) == 0) {
            val byteIndex = writePointer shr 3
            ensureCapacity(byteIndex + value.size - 1)
            System.arraycopy(value, 0, stream, byteIndex, value.size)
            writePointer += value.size shl 3
            return
        }

        value.forEach { writeUInt8(it.toInt() and 0xFF) }
    }

    fun readBytes(): ByteArray {
        val lenLong = readUInt32()
        if (lenLong == 0L) return byteArrayOf()
        if (lenLong > Int.MAX_VALUE.toLong()) {
            throw IllegalStateException("Byte array too large: $lenLong")
        }

        val length = lenLong.toInt()
        val requiredBits = length.toLong() * 8L
        if (requiredBits > remainingBits()) {
            return byteArrayOf()
        }

        val out = ByteArray(length)

        // Fast path when byte-aligned.
        if ((readPointer and 7) == 0) {
            val byteIndex = readPointer shr 3
            System.arraycopy(stream, byteIndex, out, 0, length)
            readPointer += length shl 3
            return out
        }

        for (i in 0 until length) {
            out[i] = readUInt8().toByte()
        }
        return out
    }

    private fun remainingBits(): Long = (writePointer - readPointer).toLong()

    private fun writeBits(value: Long, bits: Int) {
        if (bits <= 0) return

        reserveBits(bits.toLong())

        for (i in bits - 1 downTo 0) {
            val bit = ((value ushr i) and 1L).toInt()
            val byteIndex = writePointer shr 3
            val shift = 7 - (writePointer and 7)

            stream[byteIndex] =
                    if (bit == 1) {
                        (stream[byteIndex].toInt() or (1 shl shift)).toByte()
                    } else {
                        (stream[byteIndex].toInt() and (1 shl shift).inv()).toByte()
                    }

            writePointer++
        }
    }

    private fun readBits(bits: Int): Long {
        if (bits <= 0) return 0L
        if (remainingBits() < bits.toLong()) {
            throw IllegalStateException("Not enough bits to read")
        }

        var value = 0L
        repeat(bits) {
            val bit =
                    (stream[readPointer shr 3].toInt() ushr
                                    (7 - (readPointer and 7))) and
                            1
            value = (value shl 1) or bit.toLong()
            readPointer++
        }
        return value
    }

    private fun reserveBits(bitsToWrite: Long) {
        if (bitsToWrite <= 0L) return

        val lastBitIndex = writePointer.toLong() + bitsToWrite - 1L
        if (lastBitIndex < 0L) {
            throw IllegalStateException("Bit index overflow")
        }

        val byteIndex = lastBitIndex ushr 3
        if (byteIndex > Int.MAX_VALUE.toLong()) {
            throw IllegalStateException("Stream too large")
        }

        ensureCapacity(byteIndex.toInt())
    }

    private fun ensureCapacity(byteIndex: Int) {
        val requiredSize = byteIndex + 1
        if (requiredSize <= stream.size) return

        var newSize = if (stream.isEmpty()) 32 else stream.size
        while (newSize < requiredSize) {
            newSize = if (newSize <= Int.MAX_VALUE / 2) newSize shl 1 else requiredSize
        }

        stream = stream.copyOf(newSize)
    }
}
