#!/bin/bash

# Hardware RNG Device Testing Script
# Tests all three detected hardware RNG devices for functionality
# Author: OTP Cipher Implementation
# Version: 1.0

set -e  # Exit on any error

echo "========================================================================"
echo "Hardware RNG Device Testing Script - OTP Cipher v0.3.16"
echo "========================================================================"
echo

# Colors for output
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
NC='\033[0m' # No Color

# Test counters
TOTAL_TESTS=0
PASSED_TESTS=0
FAILED_TESTS=0

# Function to print test results
print_result() {
    local test_name="$1"
    local result="$2"
    local details="$3"
    
    TOTAL_TESTS=$((TOTAL_TESTS + 1))
    
    if [ "$result" = "PASS" ]; then
        echo -e "${GREEN}✓ PASS${NC}: $test_name"
        PASSED_TESTS=$((PASSED_TESTS + 1))
    elif [ "$result" = "FAIL" ]; then
        echo -e "${RED}✗ FAIL${NC}: $test_name"
        FAILED_TESTS=$((FAILED_TESTS + 1))
        if [ -n "$details" ]; then
            echo -e "  ${RED}Error:${NC} $details"
        fi
    elif [ "$result" = "SKIP" ]; then
        echo -e "${YELLOW}⚠ SKIP${NC}: $test_name"
        if [ -n "$details" ]; then
            echo -e "  ${YELLOW}Reason:${NC} $details"
        fi
    fi
}

# Function to test device detection
test_device_detection() {
    echo -e "${BLUE}=== Device Detection Tests ===${NC}"
    echo
    
    # Test 1: Check if devices are detected
    echo "Scanning for hardware RNG devices..."
    
    # Create a temporary test program to check device detection
    cat > temp_device_test.c << 'EOF'
#include <stdio.h>
#include <stdlib.h>
#include "include/otp.h"

int main(int argc, char* argv[]) {
    (void)argc; (void)argv; // Suppress unused parameter warnings
    
    hardware_rng_device_t devices[10];
    int num_devices_found = 0;
    
    if (detect_all_hardware_rng_devices(devices, 10, &num_devices_found) != 0) {
        printf("ERROR: Device detection failed\n");
        return 1;
    }
    
    printf("DEVICES_FOUND:%d\n", num_devices_found);
    
    for (int i = 0; i < num_devices_found; i++) {
        printf("DEVICE:%d:%s:%d:%s\n", i, devices[i].port_path, devices[i].device_type, devices[i].friendly_name);
    }
    
    return 0;
}
EOF

    # Compile the test program
    if gcc -Wall -Wextra -std=c99 -Iinclude -o temp_device_test temp_device_test.c src/trng.o src/util.o src/state.o src/pads.o src/crypto.o src/entropy.o src/ui.o nostr_chacha20.o -lm 2>/dev/null; then
        # Run device detection
        DEVICE_OUTPUT=$(./temp_device_test 2>/dev/null)
        DEVICE_COUNT=$(echo "$DEVICE_OUTPUT" | grep "DEVICES_FOUND:" | cut -d: -f2)
        
        if [ "$DEVICE_COUNT" -gt 0 ]; then
            print_result "Hardware RNG device detection" "PASS" "Found $DEVICE_COUNT devices"
            
            # Parse device information
            echo "$DEVICE_OUTPUT" | grep "DEVICE:" | while IFS=: read -r prefix index port_path device_type friendly_name; do
                echo "  Device $index: $friendly_name at $port_path (Type: $device_type)"
            done
            echo
            
            # Store device info for later tests
            echo "$DEVICE_OUTPUT" > temp_devices.txt
        else
            print_result "Hardware RNG device detection" "FAIL" "No devices found"
            echo "  Expected devices: TrueRNG, SwiftRNG variants"
            echo "  Check USB connections and device permissions"
            echo
        fi
        
        # Clean up
        rm -f temp_device_test temp_device_test.c
    else
        print_result "Device detection compilation" "FAIL" "Could not compile test program"
    fi
}

# Function to test individual device connectivity
test_device_connectivity() {
    echo -e "${BLUE}=== Device Connectivity Tests ===${NC}"
    echo
    
    if [ ! -f temp_devices.txt ]; then
        print_result "Device connectivity tests" "SKIP" "No devices detected in previous test"
        return
    fi
    
    DEVICE_COUNT=$(grep "DEVICES_FOUND:" temp_devices.txt | cut -d: -f2)
    
    if [ "$DEVICE_COUNT" -eq 0 ]; then
        print_result "Device connectivity tests" "SKIP" "No devices available for testing"
        return
    fi
    
    # Test each detected device
    grep "DEVICE:" temp_devices.txt | while IFS=: read -r prefix index port_path device_type friendly_name; do
        echo "Testing device: $friendly_name at $port_path"
        
        # Create device-specific test
        cat > temp_connectivity_test.c << 'EOF'
#include <stdio.h>
#include <stdlib.h>
#include <unistd.h>
#include <time.h>
#include "include/otp.h"

int main(int argc, char* argv[]) {
    if (argc != 4) {
        printf("Usage: %s <port_path> <device_type> <friendly_name>\n", argv[0]);
        return 1;
    }
    
    hardware_rng_device_t device;
    snprintf(device.port_path, sizeof(device.port_path), "%s", argv[1]);
    device.device_type = atoi(argv[2]);
    snprintf(device.friendly_name, sizeof(device.friendly_name), "%s", argv[3]);
    device.is_working = 0;
    
    // Test with small buffer (1KB) and short timeout
    const size_t test_bytes = 1024;
    unsigned char* test_buffer = malloc(test_bytes);
    if (!test_buffer) {
        printf("ERROR: Cannot allocate test buffer\n");
        return 1;
    }
    
    size_t collected = 0;
    time_t start_time = time(NULL);
    
    printf("TESTING:%s\n", device.friendly_name);
    
    // Test device connectivity with timeout
    int result = collect_truerng_entropy_from_device(&device, test_buffer, test_bytes, &collected, 0);
    
    time_t end_time = time(NULL);
    double test_duration = difftime(end_time, start_time);
    
    if (result == 0 && collected > 0) {
        double speed_kbps = (collected / 1024.0) / (test_duration > 0 ? test_duration : 1.0);
        printf("SUCCESS:%zu:%0.2f:%0.2f\n", collected, test_duration, speed_kbps);
    } else {
        printf("FAILED:%d:%zu:%0.2f\n", result, collected, test_duration);
    }
    
    free(test_buffer);
    return 0;
}
EOF

        # Compile and run connectivity test
        if gcc -Wall -Wextra -std=c99 -Iinclude -o temp_connectivity_test temp_connectivity_test.c src/trng.o src/util.o src/state.o src/pads.o src/crypto.o src/entropy.o src/ui.o nostr_chacha20.o -lm 2>/dev/null; then
            
            # Run test with timeout to prevent hanging
            CONNECTIVITY_OUTPUT=$(timeout 30s ./temp_connectivity_test "$port_path" "$device_type" "$friendly_name" 2>/dev/null || echo "TIMEOUT")
            
            if echo "$CONNECTIVITY_OUTPUT" | grep -q "SUCCESS:"; then
                # Parse success output
                SUCCESS_LINE=$(echo "$CONNECTIVITY_OUTPUT" | grep "SUCCESS:")
                COLLECTED=$(echo "$SUCCESS_LINE" | cut -d: -f2)
                DURATION=$(echo "$SUCCESS_LINE" | cut -d: -f3)
                SPEED=$(echo "$SUCCESS_LINE" | cut -d: -f4)
                
                print_result "Device connectivity: $friendly_name" "PASS" "Collected ${COLLECTED} bytes in ${DURATION}s (${SPEED} KB/s)"
                
            elif echo "$CONNECTIVITY_OUTPUT" | grep -q "FAILED:"; then
                # Parse failure output
                FAILED_LINE=$(echo "$CONNECTIVITY_OUTPUT" | grep "FAILED:")
                ERROR_CODE=$(echo "$FAILED_LINE" | cut -d: -f2)
                COLLECTED=$(echo "$FAILED_LINE" | cut -d: -f3)
                
                print_result "Device connectivity: $friendly_name" "FAIL" "Error code $ERROR_CODE, collected $COLLECTED bytes"
                
            elif echo "$CONNECTIVITY_OUTPUT" | grep -q "TIMEOUT"; then
                print_result "Device connectivity: $friendly_name" "FAIL" "Test timed out after 30 seconds"
                
            else
                print_result "Device connectivity: $friendly_name" "FAIL" "Unexpected test output"
            fi
            
        else
            print_result "Device connectivity test compilation: $friendly_name" "FAIL" "Could not compile test program"
        fi
        
        echo
    done
    
    # Clean up
    rm -f temp_connectivity_test temp_connectivity_test.c
}

# Function to test device configuration
test_device_configuration() {
    echo -e "${BLUE}=== Device Configuration Tests ===${NC}"
    echo
    
    if [ ! -f temp_devices.txt ]; then
        print_result "Device configuration tests" "SKIP" "No devices detected"
        return
    fi
    
    # Test serial port configuration for each device type
    grep "DEVICE:" temp_devices.txt | while IFS=: read -r prefix index port_path device_type friendly_name; do
        echo "Testing serial configuration for: $friendly_name"
        
        # Create configuration test
        cat > temp_config_test.c << 'EOF'
#include <stdio.h>
#include <stdlib.h>
#include <fcntl.h>
#include <unistd.h>
#include <termios.h>
#include "include/otp.h"

int main(int argc, char* argv[]) {
    if (argc != 3) {
        printf("Usage: %s <port_path> <device_type>\n", argv[0]);
        return 1;
    }
    
    const char* port_path = argv[1];
    hardware_rng_device_type_t device_type = atoi(argv[2]);
    
    printf("TESTING_CONFIG:%s:%d\n", port_path, device_type);
    
    // Test opening the device
    int fd = open(port_path, O_RDONLY | O_NOCTTY | O_NONBLOCK);
    if (fd < 0) {
        printf("FAILED:Cannot open device\n");
        return 1;
    }
    
    // Test configuring the serial port
    int config_result = configure_rng_serial_port(fd, device_type);
    
    if (config_result == 0) {
        printf("SUCCESS:Serial port configured successfully\n");
    } else {
        printf("FAILED:Serial port configuration failed\n");
    }
    
    close(fd);
    return 0;
}
EOF

        # Compile and run configuration test
        if gcc -Wall -Wextra -std=c99 -Iinclude -o temp_config_test temp_config_test.c src/trng.o src/util.o src/state.o src/ui.o -lm 2>/dev/null; then
            
            CONFIG_OUTPUT=$(./temp_config_test "$port_path" "$device_type" 2>/dev/null || echo "ERROR")
            
            if echo "$CONFIG_OUTPUT" | grep -q "SUCCESS:"; then
                print_result "Serial configuration: $friendly_name" "PASS" "Port configured successfully"
            elif echo "$CONFIG_OUTPUT" | grep -q "FAILED:"; then
                REASON=$(echo "$CONFIG_OUTPUT" | grep "FAILED:" | cut -d: -f2)
                print_result "Serial configuration: $friendly_name" "FAIL" "$REASON"
            else
                print_result "Serial configuration: $friendly_name" "FAIL" "Unexpected configuration result"
            fi
            
        else
            print_result "Configuration test compilation: $friendly_name" "FAIL" "Could not compile test program"
        fi
    done
    
    echo
    # Clean up
    rm -f temp_config_test temp_config_test.c
}

# Function to test entropy quality
test_entropy_quality() {
    echo -e "${BLUE}=== Entropy Quality Tests ===${NC}"
    echo
    
    if [ ! -f temp_devices.txt ]; then
        print_result "Entropy quality tests" "SKIP" "No devices detected"
        return
    fi
    
    # Test entropy quality for working devices
    grep "DEVICE:" temp_devices.txt | while IFS=: read -r prefix index port_path device_type friendly_name; do
        echo "Testing entropy quality for: $friendly_name"
        
        # Create entropy quality test
        cat > temp_quality_test.c << 'EOF'
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <time.h>
#include "include/otp.h"

// Simple entropy quality check
double calculate_byte_entropy(unsigned char* data, size_t length) {
    int counts[256] = {0};
    double entropy = 0.0;
    
    // Count byte frequencies
    for (size_t i = 0; i < length; i++) {
        counts[data[i]]++;
    }
    
    // Calculate Shannon entropy
    for (int i = 0; i < 256; i++) {
        if (counts[i] > 0) {
            double p = (double)counts[i] / length;
            entropy -= p * log2(p);
        }
    }
    
    return entropy;
}

int main(int argc, char* argv[]) {
    if (argc != 4) {
        printf("Usage: %s <port_path> <device_type> <friendly_name>\n", argv[0]);
        return 1;
    }
    
    hardware_rng_device_t device;
    snprintf(device.port_path, sizeof(device.port_path), "%s", argv[1]);
    device.device_type = atoi(argv[2]);
    snprintf(device.friendly_name, sizeof(device.friendly_name), "%s", argv[3]);
    
    // Test with 4KB sample
    const size_t test_bytes = 4096;
    unsigned char* test_buffer = malloc(test_bytes);
    if (!test_buffer) {
        printf("ERROR: Cannot allocate test buffer\n");
        return 1;
    }
    
    size_t collected = 0;
    printf("TESTING_QUALITY:%s\n", device.friendly_name);
    
    // Collect entropy sample
    int result = collect_truerng_entropy_from_device(&device, test_buffer, test_bytes, &collected, 0);
    
    if (result == 0 && collected >= 1024) {  // Need at least 1KB for meaningful analysis
        double entropy = calculate_byte_entropy(test_buffer, collected);
        double entropy_percentage = (entropy / 8.0) * 100.0;  // Max entropy is 8 bits per byte
        
        if (entropy_percentage >= 95.0) {
            printf("EXCELLENT:%0.2f:%zu\n", entropy_percentage, collected);
        } else if (entropy_percentage >= 85.0) {
            printf("GOOD:%0.2f:%zu\n", entropy_percentage, collected);
        } else if (entropy_percentage >= 70.0) {
            printf("FAIR:%0.2f:%zu\n", entropy_percentage, collected);
        } else {
            printf("POOR:%0.2f:%zu\n", entropy_percentage, collected);
        }
    } else {
        printf("FAILED:%d:%zu\n", result, collected);
    }
    
    free(test_buffer);
    return 0;
}
EOF

        # Compile and run quality test
        if gcc -Wall -Wextra -std=c99 -Iinclude -o temp_quality_test temp_quality_test.c src/trng.o src/util.o src/state.o src/ui.o -lm 2>/dev/null; then
            
            QUALITY_OUTPUT=$(timeout 30s ./temp_quality_test "$port_path" "$device_type" "$friendly_name" 2>/dev/null || echo "TIMEOUT")
            
            if echo "$QUALITY_OUTPUT" | grep -q "EXCELLENT:"; then
                QUALITY_LINE=$(echo "$QUALITY_OUTPUT" | grep "EXCELLENT:")
                PERCENTAGE=$(echo "$QUALITY_LINE" | cut -d: -f2)
                BYTES=$(echo "$QUALITY_LINE" | cut -d: -f3)
                print_result "Entropy quality: $friendly_name" "PASS" "Excellent quality (${PERCENTAGE}% entropy, ${BYTES} bytes)"
                
            elif echo "$QUALITY_OUTPUT" | grep -q "GOOD:"; then
                QUALITY_LINE=$(echo "$QUALITY_OUTPUT" | grep "GOOD:")
                PERCENTAGE=$(echo "$QUALITY_LINE" | cut -d: -f2)
                BYTES=$(echo "$QUALITY_LINE" | cut -d: -f3)
                print_result "Entropy quality: $friendly_name" "PASS" "Good quality (${PERCENTAGE}% entropy, ${BYTES} bytes)"
                
            elif echo "$QUALITY_OUTPUT" | grep -q "FAIR:"; then
                QUALITY_LINE=$(echo "$QUALITY_OUTPUT" | grep "FAIR:")
                PERCENTAGE=$(echo "$QUALITY_LINE" | cut -d: -f2)
                BYTES=$(echo "$QUALITY_LINE" | cut -d: -f3)
                print_result "Entropy quality: $friendly_name" "PASS" "Fair quality (${PERCENTAGE}% entropy, ${BYTES} bytes)"
                
            elif echo "$QUALITY_OUTPUT" | grep -q "POOR:"; then
                QUALITY_LINE=$(echo "$QUALITY_OUTPUT" | grep "POOR:")
                PERCENTAGE=$(echo "$QUALITY_LINE" | cut -d: -f2)
                BYTES=$(echo "$QUALITY_LINE" | cut -d: -f3)
                print_result "Entropy quality: $friendly_name" "FAIL" "Poor quality (${PERCENTAGE}% entropy, ${BYTES} bytes)"
                
            elif echo "$QUALITY_OUTPUT" | grep -q "FAILED:"; then
                print_result "Entropy quality: $friendly_name" "FAIL" "Could not collect sufficient entropy for analysis"
                
            elif echo "$QUALITY_OUTPUT" | grep -q "TIMEOUT"; then
                print_result "Entropy quality: $friendly_name" "FAIL" "Quality test timed out"
                
            else
                print_result "Entropy quality: $friendly_name" "FAIL" "Unexpected quality test output"
            fi
            
        else
            print_result "Quality test compilation: $friendly_name" "FAIL" "Could not compile test program"
        fi
    done
    
    echo
    # Clean up
    rm -f temp_quality_test temp_quality_test.c
}

# Function to print final summary
print_summary() {
    echo
    echo "========================================================================"
    echo -e "${BLUE}Test Summary${NC}"
    echo "========================================================================"
    echo "Total tests run: $TOTAL_TESTS"
    echo -e "Passed: ${GREEN}$PASSED_TESTS${NC}"
    echo -e "Failed: ${RED}$FAILED_TESTS${NC}"
    echo -e "Success rate: $(( (PASSED_TESTS * 100) / (TOTAL_TESTS > 0 ? TOTAL_TESTS : 1) ))%"
    echo
    
    if [ $FAILED_TESTS -eq 0 ]; then
        echo -e "${GREEN}🎉 All tests passed! Hardware RNG devices are working correctly.${NC}"
    elif [ $PASSED_TESTS -gt $FAILED_TESTS ]; then
        echo -e "${YELLOW}⚠ Some tests failed, but most devices are working.${NC}"
        echo "Check the failed tests above for specific device issues."
    else
        echo -e "${RED}❌ Multiple test failures detected.${NC}"
        echo "Hardware RNG devices may have connectivity or configuration issues."
    fi
    echo
}

# Main test execution
main() {
    echo "Starting comprehensive hardware RNG device testing..."
    echo "This will test device detection, connectivity, configuration, and entropy quality."
    echo
    
    # Ensure the OTP binary exists
    if [ ! -f "./otp" ]; then
        echo -e "${RED}Error: OTP binary not found. Please run 'make' first.${NC}"
        exit 1
    fi
    
    # Run all test suites
    test_device_detection
    test_device_connectivity
    test_device_configuration
    test_entropy_quality
    
    # Clean up temporary files
    rm -f temp_devices.txt
    
    # Print final summary
    print_summary
}

# Run main function
main "$@"