import Foundation

/// Message padding for metadata privacy — hides plaintext length.
/// Matches Python: crypto_utils.py pad_plaintext / unpad_plaintext
enum MessagePadding {

    /// Magic byte prefix to distinguish padded from legacy unpadded messages.
    private static let padMagic: UInt8 = 0x01

    /// Bucket sizes for length hiding (64B to 64KB).
    private static let padBuckets = [64, 128, 256, 512, 1024, 2048, 4096, 8192, 16384, 32768, 65536]

    /// Pad plaintext to nearest bucket size to hide message length.
    ///
    /// Format: `0x01 + plaintext + random_padding + pad_length(4B big-endian)`
    /// Prefix 0x01 distinguishes padded messages from legacy unpadded (which start with '{').
    static func pad(_ plaintext: Data) -> Data {
        var content = Data([padMagic])
        content.append(plaintext)

        // +4 for the length suffix
        let minSize = content.count + 4
        let target = padBuckets.first(where: { $0 >= minSize }) ?? minSize
        let padLen = target - content.count

        // random_padding (padLen - 4 bytes) + pad_length (4 bytes big-endian)
        var result = content
        result.append(Data.randomBytes(padLen - 4))
        result.append(UInt32(padLen).bigEndianData)
        return result
    }

    /// Remove padding. Returns raw plaintext for both padded and legacy unpadded messages.
    static func unpad(_ data: Data) -> Data {
        guard !data.isEmpty else { return data }

        // Legacy unpadded message (starts with '{' for JSON)
        guard data[data.startIndex] == padMagic else { return data }

        // Too short to be validly padded (magic + at least 4 bytes for length)
        guard data.count >= 5 else { return data }

        // Read pad_length from last 4 bytes (big-endian UInt32)
        let padLenOffset = data.count - 4
        let padLen = data.withUnsafeBytes { ptr -> UInt32 in
            var value: UInt32 = 0
            withUnsafeMutableBytes(of: &value) { dest in
                dest.copyBytes(from: UnsafeRawBufferPointer(rebasing: ptr[padLenOffset...]))
            }
            return UInt32(bigEndian: value)
        }

        // Validate padding metadata
        guard padLen >= 4, padLen <= data.count - 1 else { return data }

        // Strip: skip magic byte (index 0), take up to (data.count - padLen)
        return data[data.startIndex + 1 ..< data.startIndex + data.count - Int(padLen)]
    }
}