import XCTest
@testable import bitchat

class MessagePaddingTests: XCTestCase {
    
    func testBasicPadding() {
        let originalData = Data("Hello".utf8)
        let targetSize = 256
        
        let padded = MessagePadding.pad(originalData, toSize: targetSize)
        XCTAssertEqual(padded.count, targetSize)
        
        let unpadded = MessagePadding.unpad(padded)
        XCTAssertEqual(unpadded, originalData)
    }
    
    func testMultipleBlockSizes() {
        let testMessages = [
            "Hi",
            "This is a longer message",
            "This is an even longer message that should require a larger block size",
            String(repeating: "A", count: 500)
        ]
        
        for message in testMessages {
            let data = Data(message.utf8)
            let blockSize = MessagePadding.optimalBlockSize(for: data.count)
            
            // Block size should be reasonable
            XCTAssertGreaterThan(blockSize, data.count)
            XCTAssertTrue(MessagePadding.blockSizes.contains(blockSize) || blockSize == data.count)
            
            let padded = MessagePadding.pad(data, toSize: blockSize)
            let unpadded = MessagePadding.unpad(padded)
            
            XCTAssertEqual(unpadded, data)
        }
    }
    
    func testPaddingWithLargeData() {
        let largeData = Data(repeating: 0xFF, count: 1500)
        let blockSize = MessagePadding.optimalBlockSize(for: largeData.count)
        
        // Should use 2048 block
        XCTAssertEqual(blockSize, 2048)
        
        let padded = MessagePadding.pad(largeData, toSize: blockSize)
        XCTAssertEqual(padded.count, blockSize)
        
        let unpadded = MessagePadding.unpad(padded)
        XCTAssertEqual(unpadded, largeData)
    }
    
    func testInvalidPadding() {
        // Test empty data
        let empty = Data()
        let unpaddedEmpty = MessagePadding.unpad(empty)
        XCTAssertEqual(unpaddedEmpty, empty)
        
        // Test data with invalid padding length
        var invalidPadding = Data(repeating: 0x00, count: 100)
        invalidPadding[99] = 255 // Invalid padding length
        let result = MessagePadding.unpad(invalidPadding)
        XCTAssertEqual(result, invalidPadding) // Should return original if invalid
    }
    
    func testPaddingRandomness() {
        // Ensure padding bytes are random (not predictable)
        let data = Data("Test".utf8)
        let padded1 = MessagePadding.pad(data, toSize: 256)
        let padded2 = MessagePadding.pad(data, toSize: 256)
        
        // Same size
        XCTAssertEqual(padded1.count, padded2.count)
        
        // But different padding bytes (with very high probability)
        XCTAssertNotEqual(padded1, padded2)
        
        // Both should unpad to same data
        XCTAssertEqual(MessagePadding.unpad(padded1), data)
        XCTAssertEqual(MessagePadding.unpad(padded2), data)
    }
}