package com.kecalek.chat.crypto

import java.nio.ByteBuffer
import java.security.SecureRandom

/**
 * Bucket-based message padding for metadata privacy.
 * Pads messages to fixed bucket sizes to prevent message-length analysis.
 *
 * Format: 0x01 + plaintext + random_padding + pad_length(4 bytes big-endian)
 * Compatible with Python pad_plaintext/unpad_plaintext.
 */
object MessagePadding {

    private const val PAD_MAGIC: Byte = 0x01
    private val PAD_BUCKETS = intArrayOf(64, 128, 256, 512, 1024, 2048, 4096, 8192, 16384, 32768, 65536)
    private val secureRandom = SecureRandom()

    /**
     * Pad plaintext to nearest bucket size.
     * @param plaintext raw plaintext bytes
     * @return padded bytes: 0x01 + plaintext + random_padding + pad_length(4B)
     */
    fun pad(plaintext: ByteArray): ByteArray {
        // content = magic + plaintext
        val content = ByteArray(1 + plaintext.size)
        content[0] = PAD_MAGIC
        System.arraycopy(plaintext, 0, content, 1, plaintext.size)

        // minimum total size = content + 4 bytes for pad_length
        val minSize = content.size + 4

        // find nearest bucket
        var targetSize = minSize
        for (bucket in PAD_BUCKETS) {
            if (bucket >= minSize) {
                targetSize = bucket
                break
            }
        }

        // pad_length includes itself (4 bytes) + random padding bytes
        val padLength = targetSize - content.size
        val randomPadSize = padLength - 4

        val result = ByteArray(targetSize)
        System.arraycopy(content, 0, result, 0, content.size)

        // fill random padding
        if (randomPadSize > 0) {
            val randomBytes = ByteArray(randomPadSize)
            secureRandom.nextBytes(randomBytes)
            System.arraycopy(randomBytes, 0, result, content.size, randomPadSize)
        }

        // write pad_length as big-endian uint32 at the end
        val lenBytes = ByteBuffer.allocate(4).putInt(padLength).array()
        System.arraycopy(lenBytes, 0, result, targetSize - 4, 4)

        return result
    }

    /**
     * Remove padding from padded message.
     * @param data padded message bytes
     * @return original plaintext
     */
    fun unpad(data: ByteArray): ByteArray {
        // Legacy unpadded messages (JSON starting with '{')
        if (data.isEmpty() || data[0] != PAD_MAGIC) return data
        if (data.size < 5) return data

        // Read pad_length from last 4 bytes
        val padLength = ByteBuffer.wrap(data, data.size - 4, 4).int

        // Validate
        if (padLength < 4 || padLength > data.size - 1) return data

        // Strip magic prefix (1 byte) and padding (padLength bytes from end)
        return data.copyOfRange(1, data.size - padLength)
    }
}