#!/bin/bash
# Test script for message padding implementation

set -e

echo "=== Testing Message Padding Implementation ==="
echo ""

# Colors
GREEN='\033[0;32m'
RED='\033[0;31m'
NC='\033[0m' # No Color

# Test counter
TESTS_PASSED=0
TESTS_FAILED=0

# Function to run a test
run_test() {
    local test_name="$1"
    local test_command="$2"
    
    echo -n "Testing: $test_name... "
    if eval "$test_command" > /dev/null 2>&1; then
        echo -e "${GREEN}PASS${NC}"
        ((TESTS_PASSED++))
        return 0
    else
        echo -e "${RED}FAIL${NC}"
        ((TESTS_FAILED++))
        return 1
    fi
}

# Create a small test pad if it doesn't exist
if [ ! -f "pads/"*.pad ]; then
    echo "Creating test pad (1MB)..."
    ./build/otp-x86_64 generate 1MB
    echo ""
fi

# Get the first pad checksum
PAD_CHKSUM=$(ls pads/*.pad | head -n 1 | xargs basename | sed 's/.pad$//')

echo "Using pad: ${PAD_CHKSUM:0:16}..."
echo ""

# Test 1: Encrypt and decrypt a short message (should be padded to 256 bytes)
echo "Test 1: Short message (10 bytes -> 256 bytes padded)"
TEST_MSG="Hello Test"
ENCRYPTED=$(echo "$TEST_MSG" | ./build/otp-x86_64 encrypt ${PAD_CHKSUM:0:8})
DECRYPTED=$(echo "$ENCRYPTED" | ./build/otp-x86_64 decrypt)

if [ "$DECRYPTED" = "$TEST_MSG" ]; then
    echo -e "${GREEN}✓ Short message encryption/decryption successful${NC}"
    ((TESTS_PASSED++))
else
    echo -e "${RED}✗ Short message failed${NC}"
    echo "  Expected: $TEST_MSG"
    echo "  Got: $DECRYPTED"
    ((TESTS_FAILED++))
fi
echo ""

# Test 2: Encrypt and decrypt a medium message (should be padded to 512 bytes)
echo "Test 2: Medium message (~300 bytes -> 512 bytes padded)"
TEST_MSG=$(printf 'A%.0s' {1..300})
ENCRYPTED=$(echo "$TEST_MSG" | ./build/otp-x86_64 encrypt ${PAD_CHKSUM:0:8})
DECRYPTED=$(echo "$ENCRYPTED" | ./build/otp-x86_64 decrypt)

if [ "$DECRYPTED" = "$TEST_MSG" ]; then
    echo -e "${GREEN}✓ Medium message encryption/decryption successful${NC}"
    ((TESTS_PASSED++))
else
    echo -e "${RED}✗ Medium message failed${NC}"
    ((TESTS_FAILED++))
fi
echo ""

# Test 3: Encrypt and decrypt with special characters
echo "Test 3: Special characters and unicode"
TEST_MSG="Hello! @#$%^&*() 测试"
ENCRYPTED=$(echo "$TEST_MSG" | ./build/otp-x86_64 encrypt ${PAD_CHKSUM:0:8})
DECRYPTED=$(echo "$ENCRYPTED" | ./build/otp-x86_64 decrypt)

if [ "$DECRYPTED" = "$TEST_MSG" ]; then
    echo -e "${GREEN}✓ Special characters encryption/decryption successful${NC}"
    ((TESTS_PASSED++))
else
    echo -e "${RED}✗ Special characters failed${NC}"
    echo "  Expected: $TEST_MSG"
    echo "  Got: $DECRYPTED"
    ((TESTS_FAILED++))
fi
echo ""

# Test 4: File encryption/decryption with padding
echo "Test 4: File encryption/decryption"
TEST_FILE="/tmp/otp_test_file.txt"
echo "This is a test file for OTP encryption with padding." > "$TEST_FILE"

./build/otp-x86_64 -f "$TEST_FILE" ${PAD_CHKSUM:0:8} -a -o /tmp/test_encrypted.otp.asc
./build/otp-x86_64 decrypt /tmp/test_encrypted.otp.asc -o /tmp/test_decrypted.txt

if diff "$TEST_FILE" /tmp/test_decrypted.txt > /dev/null 2>&1; then
    echo -e "${GREEN}✓ File encryption/decryption successful${NC}"
    ((TESTS_PASSED++))
else
    echo -e "${RED}✗ File encryption/decryption failed${NC}"
    ((TESTS_FAILED++))
fi

# Cleanup
rm -f "$TEST_FILE" /tmp/test_encrypted.otp.asc /tmp/test_decrypted.txt
echo ""

# Summary
echo "=== Test Summary ==="
echo -e "Tests passed: ${GREEN}${TESTS_PASSED}${NC}"
echo -e "Tests failed: ${RED}${TESTS_FAILED}${NC}"
echo ""

if [ $TESTS_FAILED -eq 0 ]; then
    echo -e "${GREEN}All tests passed!${NC}"
    exit 0
else
    echo -e "${RED}Some tests failed.${NC}"
    exit 1
fi