#define _POSIX_C_SOURCE 200809L
#define _DEFAULT_SOURCE

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <stdint.h>
#include <unistd.h>
#include <sys/stat.h>
#include <sys/statvfs.h>
#include <sys/ioctl.h>
#include <dirent.h>
#include <time.h>
#include <ctype.h>
#include <termios.h>
#include <fcntl.h>
#include <math.h>
#include "nostr_chacha20.h"
#include "main.h"


// In-place pad entropy addition using Chacha20 or direct XOR
int add_entropy_to_pad(const char* pad_chksum, const unsigned char* entropy_data,
                       size_t entropy_size, int display_progress) {
    if (!pad_chksum || !entropy_data || entropy_size < 512) {
        printf("Error: Invalid entropy data or insufficient entropy\n");
        return 1;
    }

    // Get pad file path
    char pad_path[1024];
    char state_path[1024];
    get_pad_path(pad_chksum, pad_path, state_path);

    // Check if pad exists and get size
    struct stat pad_stat;
    if (stat(pad_path, &pad_stat) != 0) {
        printf("Error: Pad file not found: %s\n", pad_path);
        return 1;
    }

    uint64_t pad_size = pad_stat.st_size;

    // Determine entropy addition method based on entropy size vs pad size
    if (entropy_size >= pad_size) {
        // Use direct XOR when entropy >= pad size
        return add_entropy_direct_xor(pad_chksum, entropy_data, entropy_size, pad_size, display_progress);
    } else {
        // Use ChaCha20 when entropy < pad size
        return add_entropy_chacha20(pad_chksum, entropy_data, entropy_size, pad_size, display_progress);
    }
}

// Direct XOR entropy addition for large entropy sources
int add_entropy_direct_xor(const char* pad_chksum, const unsigned char* entropy_data,
                          size_t entropy_size, uint64_t pad_size, int display_progress) {
    // Get pad file path
    char pad_path[1024];
    char state_path[1024];
    get_pad_path(pad_chksum, pad_path, state_path);

    // Open pad file for read/write
    FILE* pad_file = fopen(pad_path, "r+b");
    if (!pad_file) {
        printf("Error: Cannot open pad file for modification: %s\n", pad_path);
        printf("Note: Pad files are read-only. Temporarily changing permissions...\n");

        // Try to make writable temporarily
        if (chmod(pad_path, S_IRUSR | S_IWUSR) != 0) {
            printf("Error: Cannot change pad file permissions\n");
            return 1;
        }

        pad_file = fopen(pad_path, "r+b");
        if (!pad_file) {
            printf("Error: Still cannot open pad file for modification\n");
            // Restore read-only
            chmod(pad_path, S_IRUSR);
            return 1;
        }
    }

    if (display_progress) {
        printf("Adding entropy to pad using direct XOR...\n");
        printf("Pad size: %.2f GB (%lu bytes)\n", (double)pad_size / (1024.0*1024.0*1024.0), pad_size);
        printf("Entropy size: %zu bytes\n", entropy_size);
    }

    // Process pad in chunks
    unsigned char buffer[64 * 1024]; // 64KB chunks
    size_t entropy_offset = 0;
    uint64_t offset = 0;
    time_t start_time = time(NULL);

    while (offset < pad_size) {
        size_t chunk_size = sizeof(buffer);
        if (pad_size - offset < chunk_size) {
            chunk_size = pad_size - offset;
        }

        // Read current pad data
        if (fread(buffer, 1, chunk_size, pad_file) != chunk_size) {
            printf("Error: Cannot read pad data at offset %lu\n", offset);
            fclose(pad_file);
            chmod(pad_path, S_IRUSR); // Restore read-only
            return 1;
        }

        // XOR with entropy data (wrap around if entropy smaller than pad)
        for (size_t i = 0; i < chunk_size; i++) {
            buffer[i] ^= entropy_data[entropy_offset % entropy_size];
            entropy_offset++;
        }

        // Seek back and write modified data
        if (fseek(pad_file, offset, SEEK_SET) != 0) {
            printf("Error: Cannot seek to offset %lu\n", offset);
            fclose(pad_file);
            chmod(pad_path, S_IRUSR);
            return 1;
        }

        if (fwrite(buffer, 1, chunk_size, pad_file) != chunk_size) {
            printf("Error: Cannot write modified pad data\n");
            fclose(pad_file);
            chmod(pad_path, S_IRUSR);
            return 1;
        }

        offset += chunk_size;

        // Show progress for large pads
        if (display_progress && offset % (64 * 1024 * 1024) == 0) { // Every 64MB
            show_progress(offset, pad_size, start_time);
        }
    }

    fclose(pad_file);

    // Restore read-only permissions
    if (chmod(pad_path, S_IRUSR) != 0) {
        printf("Warning: Cannot restore pad file to read-only\n");
    }

    if (display_progress) {
        show_progress(pad_size, pad_size, start_time);
        printf("\n✓ Entropy successfully added to pad using direct XOR\n");
        printf("✓ Pad integrity maintained\n");
        printf("✓ %zu bytes of entropy distributed across entire pad\n", entropy_size);
        printf("✓ Pad restored to read-only mode\n");

        // Update checksum after entropy addition
        printf("\n🔄 Updating pad checksum...\n");
        char new_chksum[65];
        int checksum_result = update_pad_checksum_after_entropy(pad_chksum, new_chksum);

        if (checksum_result == 0) {
            printf("✓ Pad checksum updated successfully\n");
            printf("  Old checksum: %.16s...\n", pad_chksum);
            printf("  New checksum: %.16s...\n", new_chksum);
            printf("✓ Pad files renamed to new checksum\n");

            // Pause before returning to menu to let user see the success message
            print_centered_header("Entropy Addition Complete", 1);
        } else if (checksum_result == 2) {
            printf("ℹ Checksum unchanged (unusual but not an error)\n");
        } else {
            printf("⚠ Warning: Checksum update failed (entropy was added successfully)\n");
            printf("  You may need to manually handle the checksum update\n");
            return 1; // Report error despite successful entropy addition
        }
    }

    return 0;
}

// ChaCha20 entropy addition for smaller entropy sources
int add_entropy_chacha20(const char* pad_chksum, const unsigned char* entropy_data,
                        size_t entropy_size, uint64_t pad_size, int display_progress) {
    // Derive Chacha20 key and nonce from entropy
    unsigned char key[32], nonce[12];
    if (derive_chacha20_params(entropy_data, entropy_size, key, nonce) != 0) {
        printf("Error: Failed to derive Chacha20 parameters from entropy\n");
        return 1;
    }

    // Get pad file path
    char pad_path[1024];
    char state_path[1024];
    get_pad_path(pad_chksum, pad_path, state_path);

    // Open pad file for read/write
    FILE* pad_file = fopen(pad_path, "r+b");
    if (!pad_file) {
        printf("Error: Cannot open pad file for modification: %s\n", pad_path);
        printf("Note: Pad files are read-only. Temporarily changing permissions...\n");

        // Try to make writable temporarily
        if (chmod(pad_path, S_IRUSR | S_IWUSR) != 0) {
            printf("Error: Cannot change pad file permissions\n");
            return 1;
        }

        pad_file = fopen(pad_path, "r+b");
        if (!pad_file) {
            printf("Error: Still cannot open pad file for modification\n");
            // Restore read-only
            chmod(pad_path, S_IRUSR);
            return 1;
        }
    }

    if (display_progress) {
        printf("Adding entropy to pad using Chacha20...\n");
        printf("Pad size: %.2f GB (%lu bytes)\n", (double)pad_size / (1024.0*1024.0*1024.0), pad_size);
    }

    // Process pad in chunks
    unsigned char buffer[64 * 1024]; // 64KB chunks
    unsigned char keystream[64 * 1024];
    uint64_t offset = 0;
    uint32_t counter = 0;
    time_t start_time = time(NULL);

    while (offset < pad_size) {
        size_t chunk_size = sizeof(buffer);
        if (pad_size - offset < chunk_size) {
            chunk_size = pad_size - offset;
        }

        // Read current pad data
        if (fread(buffer, 1, chunk_size, pad_file) != chunk_size) {
            printf("Error: Cannot read pad data at offset %lu\n", offset);
            fclose(pad_file);
            chmod(pad_path, S_IRUSR); // Restore read-only
            return 1;
        }

        // Generate keystream for this chunk
        if (chacha20_encrypt(key, counter, nonce, buffer, keystream, chunk_size) != 0) {
            printf("Error: Chacha20 keystream generation failed\n");
            fclose(pad_file);
            chmod(pad_path, S_IRUSR);
            return 1;
        }

        // XOR existing pad with keystream (adds entropy)
        for (size_t i = 0; i < chunk_size; i++) {
            buffer[i] ^= keystream[i];
        }

        // Seek back and write modified data
        if (fseek(pad_file, offset, SEEK_SET) != 0) {
            printf("Error: Cannot seek to offset %lu\n", offset);
            fclose(pad_file);
            chmod(pad_path, S_IRUSR);
            return 1;
        }

        if (fwrite(buffer, 1, chunk_size, pad_file) != chunk_size) {
            printf("Error: Cannot write modified pad data\n");
            fclose(pad_file);
            chmod(pad_path, S_IRUSR);
            return 1;
        }

        offset += chunk_size;
        counter += (chunk_size + 63) / 64; // Round up for block count

        // Show progress for large pads
        if (display_progress && offset % (64 * 1024 * 1024) == 0) { // Every 64MB
            show_progress(offset, pad_size, start_time);
        }
    }

    fclose(pad_file);

    // Restore read-only permissions
    if (chmod(pad_path, S_IRUSR) != 0) {
        printf("Warning: Cannot restore pad file to read-only\n");
    }

    if (display_progress) {
        show_progress(pad_size, pad_size, start_time);
        printf("\n✓ Entropy successfully added to pad using Chacha20\n");
        printf("✓ Pad integrity maintained\n");
        printf("✓ %zu bytes of entropy distributed across entire pad\n", entropy_size);
        printf("✓ Pad restored to read-only mode\n");

        // Update checksum after entropy addition
        printf("\n🔄 Updating pad checksum...\n");
        char new_chksum[65];
        int checksum_result = update_pad_checksum_after_entropy(pad_chksum, new_chksum);

        if (checksum_result == 0) {
            printf("✓ Pad checksum updated successfully\n");
            printf("  Old checksum: %.16s...\n", pad_chksum);
            printf("  New checksum: %.16s...\n", new_chksum);
            printf("✓ Pad files renamed to new checksum\n");

            // Pause before returning to menu to let user see the success message
            print_centered_header("Entropy Addition Complete", 1);
        } else if (checksum_result == 2) {
            printf("ℹ Checksum unchanged (unusual but not an error)\n");
        } else {
            printf("⚠ Warning: Checksum update failed (entropy was added successfully)\n");
            printf("  You may need to manually handle the checksum update\n");
            return 1; // Report error despite successful entropy addition
        }
    }

    return 0;
}

// Enhanced entropy collection with visual feedback
int collect_entropy_with_feedback(unsigned char* entropy_buffer, size_t target_bytes,
                                 size_t* collected_bytes, int allow_early_exit) {
    struct termios original_termios;
    entropy_collection_state_t state = {0};

    // Initialize state
    state.target_bytes = target_bytes;
    state.auto_complete_enabled = allow_early_exit;
    state.collection_start_time = get_precise_time();

    // Setup raw terminal
    if (setup_raw_terminal(&original_termios) != 0) {
        printf("Error: Cannot setup terminal for entropy collection\n");
        return 1;
    }

    // Clear screen area for display
    printf("\n\n\n\n\n\n");
    printf("\033[2J\033[H"); // Clear screen and move to top

    unsigned char entropy_block[16];
    struct timespec timestamp;
    uint32_t sequence_counter = 0;
    char key;
    unsigned char seen_keys[256] = {0};

    *collected_bytes = 0;

    while (state.collected_bytes < target_bytes) {
        // Update display
        state.quality_score = calculate_overall_quality(&state);
        display_entropy_progress(&state);

        // Non-blocking read
        if (read(STDIN_FILENO, &key, 1) == 1) {
            // Handle ESC key for early exit
            if (key == 27 && allow_early_exit && state.collected_bytes >= 1024) {
                break; // Early exit allowed
            }

            // Record keypress timing
            double current_time = get_precise_time();
            state.last_keypress_time = current_time;

            // Update key histogram
            state.key_histogram[(unsigned char)key]++;

            // Get high precision timestamp
            clock_gettime(CLOCK_MONOTONIC, &timestamp);

            // Create enhanced entropy block: [key][timestamp][sequence][quality_bits]
            entropy_block[0] = key;
            memcpy(&entropy_block[1], &timestamp.tv_sec, 8);
            memcpy(&entropy_block[9], &timestamp.tv_nsec, 4);
            memcpy(&entropy_block[13], &sequence_counter, 2);
            entropy_block[15] = (unsigned char)(current_time * 1000) & 0xFF; // Sub-millisecond timing

            // Add to entropy buffer
            if (state.collected_bytes + 16 <= MAX_ENTROPY_BUFFER) {
                memcpy(entropy_buffer + state.collected_bytes, entropy_block, 16);
                state.collected_bytes += 16;
            }

            sequence_counter++;

            // Track unique keys
            if (!seen_keys[(unsigned char)key]) {
                seen_keys[(unsigned char)key] = 1;
                state.unique_keys++;
            }
        } else {
            // No key available, just sleep and wait for keystrokes
            usleep(10000); // 10ms delay - wait for keystrokes, don't add timing entropy
        }

        // Auto-complete at target if enabled
        if (state.collected_bytes >= target_bytes) {
            break;
        }
    }

    // Final display update
    state.quality_score = calculate_overall_quality(&state);
    display_entropy_progress(&state);

    // Summary
    double collection_time = get_precise_time() - state.collection_start_time;
    printf("\n\n✓ Entropy collection complete!\n");
    printf("  Collected: %zu bytes in %.1f seconds\n", state.collected_bytes, collection_time);
    printf("  Quality: %d%% (Excellent: 80%%+, Good: 60%%+)\n", state.quality_score);
    printf("  Unique keys: %zu\n", state.unique_keys);

    // Restore terminal
    restore_terminal(&original_termios);

    *collected_bytes = state.collected_bytes;
    return 0;
}

// Chacha20 key derivation from collected entropy
int derive_chacha20_params(const unsigned char* entropy_data, size_t entropy_size,
                          unsigned char key[32], unsigned char nonce[12]) {
    if (!entropy_data || entropy_size < 512 || !key || !nonce) {
        return 1; // Error: insufficient entropy or null pointers
    }

    // Phase 1: Generate base key from entropy using enhanced XOR checksum method
    unsigned char enhanced_checksum[44]; // 32 key + 12 nonce
    memset(enhanced_checksum, 0, 44);

    // Mix entropy data similar to calculate_checksum but for 44 bytes
    for (size_t i = 0; i < entropy_size; i++) {
        unsigned char bucket = i % 44;
        enhanced_checksum[bucket] ^= entropy_data[i] ^
                                   ((i >> 8) & 0xFF) ^
                                   ((i >> 16) & 0xFF) ^
                                   ((i >> 24) & 0xFF);
    }

    // Phase 2: Add system entropy for additional randomness
    unsigned char system_entropy[32];
    FILE* urandom = fopen("/dev/urandom", "rb");
    if (!urandom) {
        return 2; // Error: cannot access system entropy
    }

    if (fread(system_entropy, 1, 32, urandom) != 32) {
        fclose(urandom);
        return 2; // Error: insufficient system entropy
    }
    fclose(urandom);

    // Mix system entropy into derived key
    for (int i = 0; i < 32; i++) {
        enhanced_checksum[i] ^= system_entropy[i];
    }

    // Extract key and nonce
    memcpy(key, enhanced_checksum, 32);
    memcpy(nonce, enhanced_checksum + 32, 12);

    return 0; // Success
}

// Get file path and size information for entropy collection
int get_file_entropy_info(char* file_path, size_t max_path_len, size_t* file_size, int display_progress) {
    if (display_progress) {
        print_centered_header("File Entropy Collection", 0);
        printf("Load entropy from binary file (.bin format)\n");
    }

    printf("Enter path to binary entropy file: ");
    fflush(stdout);

    if (!fgets(file_path, max_path_len, stdin)) {
        printf("Error: Failed to read input\n");
        return 1;
    }

    // Remove newline
    file_path[strcspn(file_path, "\n")] = 0;

    // Check if file exists and get size
    struct stat file_stat;
    if (stat(file_path, &file_stat) != 0) {
        printf("Error: File '%s' not found\n", file_path);
        return 1;
    }

    if (!S_ISREG(file_stat.st_mode)) {
        printf("Error: '%s' is not a regular file\n", file_path);
        return 1;
    }

    *file_size = file_stat.st_size;
    if (*file_size == 0) {
        printf("Error: File is empty\n");
        return 1;
    }

    if (display_progress) {
        printf("✓ File found: %s\n", file_path);
        printf("  Size: %zu bytes\n", *file_size);
    }

    return 0; // Success
}

// Collect entropy from binary file (legacy function for backward compatibility)
int collect_file_entropy(unsigned char* entropy_buffer, size_t target_bytes,
                        size_t* collected_bytes, int display_progress) {
    char file_path[512];
    size_t file_size;

    // Get file path and size first
    if (get_file_entropy_info(file_path, sizeof(file_path), &file_size, display_progress) != 0) {
        return 1;
    }

    if (file_size < target_bytes) {
        printf("Warning: File size (%zu bytes) is smaller than target (%zu bytes)\n",
               file_size, target_bytes);
        printf("Will read available data and pad with zeros if necessary.\n");
    }

    // Open file for reading
    FILE* entropy_file = fopen(file_path, "rb");
    if (!entropy_file) {
        printf("Error: Cannot open file '%s' for reading\n", file_path);
        return 1;
    }

    if (display_progress) {
        printf("Reading entropy from file...\n");
    }

    // Read entropy data
    size_t bytes_to_read = (file_size < target_bytes) ? file_size : target_bytes;
    size_t bytes_read = fread(entropy_buffer, 1, bytes_to_read, entropy_file);

    if (bytes_read != bytes_to_read) {
        printf("Error: Failed to read %zu bytes from file (read %zu)\n",
               bytes_to_read, bytes_read);
        fclose(entropy_file);
        return 1;
    }

    fclose(entropy_file);

    // Pad with zeros if file was smaller than target
    if (bytes_read < target_bytes) {
        memset(entropy_buffer + bytes_read, 0, target_bytes - bytes_read);
        *collected_bytes = target_bytes; // We padded to target size
    } else {
        *collected_bytes = bytes_read;
    }

    if (display_progress) {
        printf("✓ File entropy collection complete!\n");
        printf("  File: %s\n", file_path);
        printf("  Read: %zu bytes\n", bytes_read);
        printf("  Total: %zu bytes (padded to target if necessary)\n", *collected_bytes);
    }

    return 0; // Success
}

// Add file entropy directly to pad using streaming (for large files)
int add_file_entropy_streaming(const char* pad_chksum, const char* file_path, size_t file_size, int display_progress) {
    // Get pad file path
    char pad_path[1024];
    char state_path[1024];
    get_pad_path(pad_chksum, pad_path, state_path);

    // Check if pad exists and get size
    struct stat pad_stat;
    if (stat(pad_path, &pad_stat) != 0) {
        printf("Error: Pad file not found: %s\n", pad_path);
        return 1;
    }

    uint64_t pad_size = pad_stat.st_size;

    // Open entropy file for reading
    FILE* entropy_file = fopen(file_path, "rb");
    if (!entropy_file) {
        printf("Error: Cannot open entropy file '%s' for reading\n", file_path);
        return 1;
    }

    // Open pad file for read/write
    FILE* pad_file = fopen(pad_path, "r+b");
    if (!pad_file) {
        printf("Error: Cannot open pad file for modification: %s\n", pad_path);
        fclose(entropy_file);
        return 1;
    }

    if (display_progress) {
        printf("Adding entropy to pad using streaming direct XOR...\n");
        printf("Pad size: %.2f GB (%lu bytes)\n", (double)pad_size / (1024.0*1024.0*1024.0), pad_size);
        printf("Entropy file: %.2f GB (%zu bytes)\n", (double)file_size / (1024.0*1024.0*1024.0), file_size);
    }

    // Process in chunks
    unsigned char pad_buffer[64 * 1024];
    unsigned char entropy_buffer[64 * 1024];
    uint64_t offset = 0;
    size_t entropy_offset = 0;
    time_t start_time = time(NULL);

    while (offset < pad_size) {
        size_t chunk_size = sizeof(pad_buffer);
        if (pad_size - offset < chunk_size) {
            chunk_size = pad_size - offset;
        }

        // Read current pad data
        if (fread(pad_buffer, 1, chunk_size, pad_file) != chunk_size) {
            printf("Error: Cannot read pad data at offset %lu\n", offset);
            fclose(entropy_file);
            fclose(pad_file);
            return 1;
        }

        // Read entropy data (wrap around if file smaller than pad)
        size_t entropy_read = 0;
        while (entropy_read < chunk_size) {
            size_t to_read = chunk_size - entropy_read;
            if (to_read > sizeof(entropy_buffer)) {
                to_read = sizeof(entropy_buffer);
            }

            size_t read_bytes = fread(entropy_buffer, 1, to_read, entropy_file);
            if (read_bytes == 0) {
                // Reached end of entropy file, wrap around
                fseek(entropy_file, 0, SEEK_SET);
                entropy_offset = 0;
                read_bytes = fread(entropy_buffer, 1, to_read, entropy_file);
                if (read_bytes == 0) {
                    printf("Error: Cannot read from entropy file\n");
                    fclose(entropy_file);
                    fclose(pad_file);
                    return 1;
                }
            }

            // XOR this chunk
            for (size_t i = 0; i < read_bytes; i++) {
                pad_buffer[entropy_read + i] ^= entropy_buffer[i];
            }

            entropy_read += read_bytes;
            entropy_offset += read_bytes;
        }

        // Seek back and write modified data
        if (fseek(pad_file, offset, SEEK_SET) != 0) {
            printf("Error: Cannot seek to offset %lu\n", offset);
            fclose(entropy_file);
            fclose(pad_file);
            return 1;
        }

        if (fwrite(pad_buffer, 1, chunk_size, pad_file) != chunk_size) {
            printf("Error: Cannot write modified pad data\n");
            fclose(entropy_file);
            fclose(pad_file);
            return 1;
        }

        offset += chunk_size;

        // Show progress for large pads
        if (display_progress && offset % (64 * 1024 * 1024) == 0) {
            show_progress(offset, pad_size, start_time);
        }
    }

    fclose(entropy_file);
    fclose(pad_file);

    if (display_progress) {
        show_progress(pad_size, pad_size, start_time);
        printf("\n✓ Entropy successfully added to pad using streaming direct XOR\n");
        printf("✓ Pad integrity maintained\n");
        printf("✓ %zu bytes of entropy distributed across entire pad\n", file_size);
    }

    return 0;
}

// Collect entropy by source type with unified interface
int collect_entropy_by_source(entropy_source_t source, unsigned char* entropy_buffer,
                              size_t target_bytes, size_t* collected_bytes, int display_progress) {
    switch (source) {
        case ENTROPY_SOURCE_KEYBOARD:
            return collect_entropy_with_feedback(entropy_buffer, target_bytes, collected_bytes, 1);

        case ENTROPY_SOURCE_TRUERNG:
            return collect_truerng_entropy(entropy_buffer, target_bytes, collected_bytes, display_progress);

        case ENTROPY_SOURCE_DICE:
            return collect_dice_entropy(entropy_buffer, target_bytes, collected_bytes, display_progress);

        case ENTROPY_SOURCE_FILE:
            return collect_file_entropy(entropy_buffer, target_bytes, collected_bytes, display_progress);

        default:
            if (display_progress) {
                printf("Error: Unknown entropy source\n");
            }
            return 1;
    }
}

// Collect manual entropy from any printable character input
int collect_dice_entropy(unsigned char* entropy_buffer, size_t target_bytes,
                        size_t* collected_bytes, int display_progress) {
    if (display_progress) {
        print_centered_header("Manual Entropy Collection", 0);
        printf("Enter any text, numbers, or symbols for entropy.\n");
        printf("Target: %zu bytes (%zu characters needed)\n", target_bytes, target_bytes);
        printf("Press Enter after each line, or 'done' when finished.\n\n");
    }

    size_t bytes_written = 0;

    char input[256];

    while (bytes_written < target_bytes) {
        if (display_progress) {
            double percentage = (double)bytes_written / target_bytes * 100.0;
            printf("Progress: %.1f%% (%zu/%zu bytes) - Enter text: ",
                   percentage, bytes_written, target_bytes);
            fflush(stdout);
        }

        if (!fgets(input, sizeof(input), stdin)) {
            if (display_progress) {
                printf("Error: Failed to read input\n");
            }
            return 1;
        }

        // Remove newline
        input[strcspn(input, "\n")] = 0;

        // Check for done command
        if (strcmp(input, "done") == 0 && bytes_written >= target_bytes / 2) {
            break; // Allow early exit if we have at least half the target
        }

        // Process each printable character as 8 bits of entropy
        for (size_t i = 0; input[i] && bytes_written < target_bytes; i++) {
            char c = input[i];
            if (c >= 32 && c <= 126) { // Printable ASCII characters
                entropy_buffer[bytes_written++] = (unsigned char)c;
            }
        }
    }

    if (display_progress) {
        printf("\n✓ Manual entropy collection complete!\n");
        printf("  Collected: %zu bytes from text input\n", bytes_written);
        printf("  Entropy quality: 8 bits per character\n");
    }

    *collected_bytes = bytes_written;
    return 0; // Success
}

void restore_terminal(struct termios* original_termios) {
    tcsetattr(STDIN_FILENO, TCSANOW, original_termios);

    // Reset stdin to blocking
    int flags = fcntl(STDIN_FILENO, F_GETFL);
    fcntl(STDIN_FILENO, F_SETFL, flags & ~O_NONBLOCK);
}
double get_precise_time(void) {
    struct timespec ts;
    clock_gettime(CLOCK_MONOTONIC, &ts);
    return ts.tv_sec + ts.tv_nsec / 1000000000.0;
}

void draw_progress_bar(double percentage, int width) {
    int filled = (int)(percentage / 100.0 * width);
    if (filled > width) filled = width;

    printf("[");
    for (int i = 0; i < filled; i++) {
        printf("█");
    }
    for (int i = filled; i < width; i++) {
        printf("░");
    }
    printf("]");
}

void draw_quality_bar(double quality, int width, const char* label) {
    int filled = (int)(quality / 100.0 * width);
    if (filled > width) filled = width;

    // Color coding based on quality
    const char* color;
    if (quality >= 80) color = "\033[32m"; // Green
    else if (quality >= 60) color = "\033[33m"; // Yellow
    else color = "\033[31m"; // Red

    printf("%s", color);
    draw_progress_bar(quality, width);
    printf("\033[0m %-10s", label); // Reset color
}

double calculate_timing_quality(const entropy_collection_state_t* state) {
    // Analyze timing variance between keypresses
    if (state->collected_bytes < 32) return 0.0; // Need minimum data

    // Simplified timing quality based on collection rate and variation
    double elapsed = get_precise_time() - state->collection_start_time;
    if (elapsed < 0.1) return 0.0;

    double rate = state->collected_bytes / elapsed;

    // Optimal rate is around 50-200 bytes/second (moderate typing with good timing variance)
    if (rate >= 50 && rate <= 200) return 90.0;
    if (rate >= 20 && rate <= 500) return 70.0;
    if (rate >= 10 && rate <= 1000) return 50.0;
    return 30.0;
}

double calculate_variety_quality(const entropy_collection_state_t* state) {
    // Analyze key variety and distribution
    if (state->collected_bytes < 16) return 0.0;

    // Calculate entropy from key histogram
    double entropy = 0.0;
    size_t total_keys = 0;

    // Count total keypresses
    for (int i = 0; i < 256; i++) {
        total_keys += state->key_histogram[i];
    }

    if (total_keys == 0) return 0.0;

    // Calculate Shannon entropy
    for (int i = 0; i < 256; i++) {
        if (state->key_histogram[i] > 0) {
            double p = (double)state->key_histogram[i] / total_keys;
            entropy -= p * log2(p);
        }
    }

    // Convert entropy to quality score (0-100)
    double max_entropy = log2(256); // Perfect entropy for 8-bit keyspace
    double normalized_entropy = entropy / max_entropy;

    // Scale based on unique keys as well
    double unique_key_factor = (double)state->unique_keys / 50.0; // 50+ unique keys is excellent
    if (unique_key_factor > 1.0) unique_key_factor = 1.0;

    return (normalized_entropy * 70.0 + unique_key_factor * 30.0);
}

unsigned char calculate_overall_quality(const entropy_collection_state_t* state) {
    double timing = calculate_timing_quality(state);
    double variety = calculate_variety_quality(state);

    // Simple collection progress bonus
    double progress_bonus = (double)state->collected_bytes / state->target_bytes * 20.0;
    if (progress_bonus > 20.0) progress_bonus = 20.0;

    // Weighted average
    double overall = (timing * 0.4 + variety * 0.4 + progress_bonus);
    if (overall > 100.0) overall = 100.0;

    return (unsigned char)overall;
}

void display_entropy_progress(const entropy_collection_state_t* state) {
    // Calculate percentages
    double progress = (double)state->collected_bytes / state->target_bytes * 100.0;
    if (progress > 100.0) progress = 100.0;

    double quality = state->quality_score;
    double timing_quality = calculate_timing_quality(state);
    double variety_quality = calculate_variety_quality(state);

    // Clear previous output and redraw
    printf("\033[2K\r"); // Clear line
    printf("\033[A\033[2K\r"); // Move up and clear
    printf("\033[A\033[2K\r"); // Move up and clear
    printf("\033[A\033[2K\r"); // Move up and clear
    printf("\033[A\033[2K\r"); // Move up and clear
    printf("\033[A\033[2K\r"); // Move up and clear

    // Header
    printf("Adding Entropy to Pad - Target: %zu bytes\n\n", state->target_bytes);

    // Main progress bar
    printf("Progress: ");
    draw_progress_bar(progress, 50);
    printf(" %.1f%% (%zu/%zu bytes)\n", progress, state->collected_bytes, state->target_bytes);

    // Quality indicators
    printf("Quality:  ");
    draw_quality_bar(quality, 50, "OVERALL");
    printf("\n");

    printf("Timing:   ");
    draw_quality_bar(timing_quality, 50, "VARIED");
    printf("\n");

    printf("Keys:     ");
    draw_quality_bar(variety_quality, 50, "DIVERSE");
    printf("\n");

    // Instructions
    if (state->collected_bytes >= 1024 && state->auto_complete_enabled) {
        printf("\nPress ESC to finish (minimum reached) or continue typing...");
    } else if (state->collected_bytes < 1024) {
        printf("\nType random keys... (%zu more bytes needed)", 1024 - state->collected_bytes);
    } else {
        printf("\nType random keys or press ESC when satisfied...");
    }

    fflush(stdout);
}

// Keyboard entropy functions
int setup_raw_terminal(struct termios* original_termios) {
    struct termios new_termios;

    if (tcgetattr(STDIN_FILENO, original_termios) != 0) {
        return 1;
    }

    new_termios = *original_termios;
    new_termios.c_lflag &= ~(ICANON | ECHO);
    new_termios.c_cc[VMIN] = 0;
    new_termios.c_cc[VTIME] = 0;

    if (tcsetattr(STDIN_FILENO, TCSANOW, &new_termios) != 0) {
        return 1;
    }

    // Set stdin to non-blocking
    int flags = fcntl(STDIN_FILENO, F_GETFL);
    if (fcntl(STDIN_FILENO, F_SETFL, flags | O_NONBLOCK) == -1) {
        tcsetattr(STDIN_FILENO, TCSANOW, original_termios);
        return 1;
    }

    return 0;
}