#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 <dirent.h>
#include <time.h>
#include <ctype.h>
#include <termios.h>
#include <fcntl.h>
#include "src/version.h"

// Custom base64 character set
static const char base64_chars[] = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/";
static const int base64_decode_table[256] = {
    -1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,
    -1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,
    -1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,62,-1,-1,-1,63,
    52,53,54,55,56,57,58,59,60,61,-1,-1,-1,-2,-1,-1,
    -1, 0, 1, 2, 3, 4, 5, 6, 7, 8, 9,10,11,12,13,14,
    15,16,17,18,19,20,21,22,23,24,25,-1,-1,-1,-1,-1,
    -1,26,27,28,29,30,31,32,33,34,35,36,37,38,39,40,
    41,42,43,44,45,46,47,48,49,50,51,-1,-1,-1,-1,-1,
    -1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,
    -1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,
    -1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,
    -1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,
    -1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,
    -1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,
    -1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,
    -1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1
};

#define MAX_INPUT_SIZE 4096
#define MAX_LINE_LENGTH 1024
#define MAX_HASH_LENGTH 65
#define PROGRESS_UPDATE_INTERVAL (64 * 1024 * 1024)  // 64MB intervals
#define PADS_DIR "pads"
#define MAX_ENTROPY_BUFFER 32768  // 32KB entropy buffer

// Function prototypes
int main(int argc, char* argv[]);
int interactive_mode(void);
int command_line_mode(int argc, char* argv[]);

// Core functions
int generate_pad(uint64_t size_bytes, int show_progress);
int generate_pad_with_entropy(uint64_t size_bytes, int show_progress, int use_keyboard_entropy);
int encrypt_text(const char* pad_identifier);
int decrypt_text(const char* pad_identifier);

// Keyboard entropy functions
int setup_raw_terminal(struct termios* original_termios);
void restore_terminal(struct termios* original_termios);
int collect_keyboard_entropy(unsigned char* entropy_buffer, size_t max_size, size_t* collected);
void simple_entropy_mix(unsigned char* urandom_buffer, size_t buffer_size, 
                         const unsigned char* entropy_data, size_t entropy_size);

// Directory management
int ensure_pads_directory(void);
void get_pad_path(const char* hash, char* pad_path, char* state_path);

// Utility functions
uint64_t parse_size_string(const char* size_str);
char* find_pad_by_prefix(const char* prefix);
int list_available_pads(void);
int show_pad_info(const char* hash);
int get_user_choice(int min, int max);
void show_progress(uint64_t current, uint64_t total, time_t start_time);

// File operations
int read_state_offset(const char* pad_hash, uint64_t* offset);
int write_state_offset(const char* pad_hash, uint64_t offset);
int calculate_checksum(const char* filename, char* checksum_hex);
void xor_checksum_256(const unsigned char* data, size_t len, unsigned char checksum[32]);
char* custom_base64_encode(const unsigned char* input, int length);
unsigned char* custom_base64_decode(const char* input, int* output_length);

// Menu functions
void show_main_menu(void);
int handle_generate_menu(void);
int handle_encrypt_menu(void);
int handle_decrypt_menu(void);

void print_usage(const char* program_name);

int main(int argc, char* argv[]) {
    if (argc == 1) {
        return interactive_mode();
    } else {
        return command_line_mode(argc, argv);
    }
}

int interactive_mode(void) {
    printf("=== OTP Cipher %s ===\n\n", get_version());
    
    while (1) {
        show_main_menu();
        int choice = get_user_choice(1, 6);
        
        switch (choice) {
            case 1:
                handle_generate_menu();
                break;
            case 2:
                handle_encrypt_menu();
                break;
            case 3:
                handle_decrypt_menu();
                break;
            case 4:
                list_available_pads();
                break;
            case 5: {
                printf("Enter pad checksum (or prefix): ");
                char input[MAX_HASH_LENGTH];
                if (fgets(input, sizeof(input), stdin)) {
                    input[strcspn(input, "\n")] = 0;
                    char* hash = find_pad_by_prefix(input);
                    if (hash) {
                        show_pad_info(hash);
                        free(hash);
                    }
                }
                break;
            }
            case 6:
                printf("Goodbye!\n");
                return 0;
        }
        printf("\n");
    }
}

int command_line_mode(int argc, char* argv[]) {
    if (strcmp(argv[1], "generate") == 0) {
        if (argc != 3) {
            printf("Usage: %s generate <size>\n", argv[0]);
            printf("Size examples: 1024, 1GB, 5TB, 512MB\n");
            return 1;
        }
        uint64_t size = parse_size_string(argv[2]);
        if (size == 0) {
            printf("Error: Invalid size format\n");
            return 1;
        }
        return generate_pad_with_entropy(size, 1, 0);  // No keyboard entropy for command line
    }
    else if (strcmp(argv[1], "encrypt") == 0) {
        if (argc != 3) {
            printf("Usage: %s encrypt <pad_hash_or_prefix>\n", argv[0]);
            return 1;
        }
        return encrypt_text(argv[2]);
    }
    else if (strcmp(argv[1], "decrypt") == 0) {
        if (argc != 3) {
            printf("Usage: %s decrypt <pad_hash_or_prefix>\n", argv[0]);
            return 1;
        }
        return decrypt_text(argv[2]);
    }
    else if (strcmp(argv[1], "list") == 0) {
        return list_available_pads();
    }
    else {
        print_usage(argv[0]);
        return 1;
    }
}

void show_main_menu(void) {
    printf("=== Main Menu ===\n");
    printf("1. Generate new pad\n");
    printf("2. Encrypt message\n");
    printf("3. Decrypt message\n");
    printf("4. List available pads\n");
    printf("5. Show pad information\n");
    printf("6. Exit\n");
    printf("\nSelect option (1-6): ");
}

int handle_generate_menu(void) {
    printf("\n=== Generate New Pad ===\n");
    printf("Enter pad size (examples: 1GB, 5TB, 512MB, 2048): ");
    
    char size_input[64];
    if (!fgets(size_input, sizeof(size_input), stdin)) {
        printf("Error: Failed to read input\n");
        return 1;
    }
    
    size_input[strcspn(size_input, "\n")] = 0;
    uint64_t size = parse_size_string(size_input);
    
    if (size == 0) {
        printf("Error: Invalid size format\n");
        return 1;
    }
    
    // Ask about keyboard entropy
    printf("\nAdd keyboard entropy for enhanced security? (y/N): ");
    char entropy_choice[10];
    int use_keyboard_entropy = 0;
    
    if (fgets(entropy_choice, sizeof(entropy_choice), stdin)) {
        if (entropy_choice[0] == 'y' || entropy_choice[0] == 'Y') {
            use_keyboard_entropy = 1;
        }
    }
    
    double size_gb = (double)size / (1024.0 * 1024.0 * 1024.0);
    if (use_keyboard_entropy) {
        printf("Generating %.2f GB pad with keyboard entropy...\n", size_gb);
    } else {
        printf("Generating %.2f GB pad...\n", size_gb);
    }
    
    return generate_pad_with_entropy(size, 1, use_keyboard_entropy);
}

int handle_encrypt_menu(void) {
    printf("\n=== Encrypt Message ===\n");
    
    int pad_count = list_available_pads();
    if (pad_count == 0) {
        printf("No pads available. Generate a pad first.\n");
        return 1;
    }
    
    printf("\nEnter pad selection (number, hash, or prefix): ");
    char input[MAX_HASH_LENGTH];
    if (!fgets(input, sizeof(input), stdin)) {
        printf("Error: Failed to read input\n");
        return 1;
    }
    
    input[strcspn(input, "\n")] = 0;
    return encrypt_text(input);
}

int handle_decrypt_menu(void) {
    printf("\n=== Decrypt Message ===\n");
    return decrypt_text(NULL);  // No pad selection needed - hash comes from message
}

uint64_t parse_size_string(const char* size_str) {
    if (!size_str) return 0;
    
    char* endptr;
    double value = strtod(size_str, &endptr);
    
    if (value <= 0) return 0;
    
    // Skip whitespace
    while (*endptr && isspace(*endptr)) endptr++;
    
    uint64_t multiplier = 1;
    
    if (*endptr) {
        char unit[4];
        strncpy(unit, endptr, 3);
        unit[3] = '\0';
        
        // Convert to uppercase
        for (int i = 0; unit[i]; i++) {
            unit[i] = toupper(unit[i]);
        }
        
        if (strcmp(unit, "K") == 0 || strcmp(unit, "KB") == 0) {
            multiplier = 1024ULL;
        } else if (strcmp(unit, "M") == 0 || strcmp(unit, "MB") == 0) {
            multiplier = 1024ULL * 1024ULL;
        } else if (strcmp(unit, "G") == 0 || strcmp(unit, "GB") == 0) {
            multiplier = 1024ULL * 1024ULL * 1024ULL;
        } else if (strcmp(unit, "T") == 0 || strcmp(unit, "TB") == 0) {
            multiplier = 1024ULL * 1024ULL * 1024ULL * 1024ULL;
        } else {
            return 0; // Invalid unit
        }
    } else {
        // No unit specified, treat as bytes
        multiplier = 1;
    }
    
    return (uint64_t)(value * multiplier);
}

char* find_pad_by_prefix(const char* prefix) {
    DIR* dir = opendir(PADS_DIR);
    if (!dir) return NULL;
    
    struct dirent* entry;
    char* matches[100];  // Store up to 100 matches
    int match_count = 0;
    
    // Check if it's a number (for interactive menu selection)
    char* endptr;
    int selection = strtol(prefix, &endptr, 10);
    if (*endptr == '\0' && selection > 0) {
        // It's a number, find the nth pad
        int current = 0;
        rewinddir(dir);
        while ((entry = readdir(dir)) != NULL && match_count == 0) {
            if (strstr(entry->d_name, ".pad") && strlen(entry->d_name) == 68) { // 64 char hash + ".pad"
                current++;
                if (current == selection) {
                    matches[match_count] = malloc(65);
                    strncpy(matches[match_count], entry->d_name, 64);
                    matches[match_count][64] = '\0';
                    match_count = 1;
                }
            }
        }
    } else {
        // Find pads that start with the prefix
        size_t prefix_len = strlen(prefix);
        while ((entry = readdir(dir)) != NULL && match_count < 100) {
            if (strstr(entry->d_name, ".pad") && strlen(entry->d_name) == 68) {
                if (strncmp(entry->d_name, prefix, prefix_len) == 0) {
                    matches[match_count] = malloc(65);
                    strncpy(matches[match_count], entry->d_name, 64);
                    matches[match_count][64] = '\0';
                    match_count++;
                }
            }
        }
    }
    
    closedir(dir);
    
    if (match_count == 0) {
        printf("No pads found matching '%s'\n", prefix);
        printf("Available pads:\n");
        list_available_pads();
        return NULL;
    } else if (match_count == 1) {
        char* result = matches[0];
        return result;
    } else {
        printf("Multiple matches found for '%s':\n", prefix);
        for (int i = 0; i < match_count; i++) {
            printf("%d. %.16s...\n", i + 1, matches[i]);
            if (i > 0) free(matches[i]);
        }
        printf("Please be more specific.\n");
        char* result = matches[0];
        for (int i = 1; i < match_count; i++) {
            free(matches[i]);
        }
        return result;
    }
}

int list_available_pads(void) {
    DIR* dir = opendir(PADS_DIR);
    if (!dir) {
        printf("Error: Cannot open pads directory\n");
        return 0;
    }
    
    struct dirent* entry;
    int count = 0;
    
    printf("Available pads:\n");
    printf("%-4s %-20s %-12s %-12s %-8s\n", "No.", "Hash (first 16 chars)", "Size", "Used", "% Used");
    printf("%-4s %-20s %-12s %-12s %-8s\n", "---", "-------------------", "----------", "----------", "------");
    
    while ((entry = readdir(dir)) != NULL) {
        if (strstr(entry->d_name, ".pad") && strlen(entry->d_name) == 68) {
            count++;
            char hash[65];
            strncpy(hash, entry->d_name, 64);
            hash[64] = '\0';
            
            // Get pad file size
            char full_path[300]; // Increased buffer size to accommodate longer paths
            snprintf(full_path, sizeof(full_path), "%s/%s", PADS_DIR, entry->d_name);
            struct stat st;
            if (stat(full_path, &st) == 0) {
                // Get used bytes from state
                uint64_t used_bytes;
                read_state_offset(hash, &used_bytes);
                
                // Format sizes
                char size_str[32], used_str[32];
                
                // Format total size
                if (st.st_size < 1024) {
                    snprintf(size_str, sizeof(size_str), "%luB", st.st_size);
                } else if (st.st_size < 1024 * 1024) {
                    snprintf(size_str, sizeof(size_str), "%.1fKB", (double)st.st_size / 1024.0);
                } else if (st.st_size < 1024 * 1024 * 1024) {
                    snprintf(size_str, sizeof(size_str), "%.1fMB", (double)st.st_size / (1024.0 * 1024.0));
                } else {
                    snprintf(size_str, sizeof(size_str), "%.2fGB", (double)st.st_size / (1024.0 * 1024.0 * 1024.0));
                }
                
                // Format used size
                if (used_bytes < 1024) {
                    snprintf(used_str, sizeof(used_str), "%luB", used_bytes);
                } else if (used_bytes < 1024 * 1024) {
                    snprintf(used_str, sizeof(used_str), "%.1fKB", (double)used_bytes / 1024.0);
                } else if (used_bytes < 1024 * 1024 * 1024) {
                    snprintf(used_str, sizeof(used_str), "%.1fMB", (double)used_bytes / (1024.0 * 1024.0));
                } else {
                    snprintf(used_str, sizeof(used_str), "%.2fGB", (double)used_bytes / (1024.0 * 1024.0 * 1024.0));
                }
                
                // Calculate percentage
                double percentage = (double)used_bytes / st.st_size * 100.0;
                
                printf("%-4d %-20.16s %-12s %-12s %.1f%%\n", count, hash, size_str, used_str, percentage);
            }
        }
    }
    
    closedir(dir);
    
    if (count == 0) {
        printf("No pads found.\n");
    }
    
    return count;
}

int show_pad_info(const char* hash) {
    char pad_filename[MAX_HASH_LENGTH + 10];
    char state_filename[MAX_HASH_LENGTH + 10];
    
    snprintf(pad_filename, sizeof(pad_filename), "%s.pad", hash);
    snprintf(state_filename, sizeof(state_filename), "%s.state", hash);
    
    struct stat st;
    if (stat(pad_filename, &st) != 0) {
        printf("Pad not found: %s\n", hash);
        return 1;
    }
    
    uint64_t used_bytes;
    read_state_offset(hash, &used_bytes);
    
    printf("=== Pad Information ===\n");
    printf("Hash: %s\n", hash);
    printf("File: %s\n", pad_filename);
    
    double size_gb = (double)st.st_size / (1024.0 * 1024.0 * 1024.0);
    double used_gb = (double)used_bytes / (1024.0 * 1024.0 * 1024.0);
    double remaining_gb = (double)(st.st_size - used_bytes) / (1024.0 * 1024.0 * 1024.0);
    
    printf("Total size: %.2f GB (%lu bytes)\n", size_gb, st.st_size);
    printf("Used: %.2f GB (%lu bytes)\n", used_gb, used_bytes);
    printf("Remaining: %.2f GB (%lu bytes)\n", remaining_gb, st.st_size - used_bytes);
    printf("Usage: %.1f%%\n", (double)used_bytes / st.st_size * 100.0);
    
    return 0;
}

int get_user_choice(int min, int max) {
    char input[64];
    int choice;
    
    while (1) {
        if (fgets(input, sizeof(input), stdin)) {
            choice = atoi(input);
            if (choice >= min && choice <= max) {
                return choice;
            }
        }
        printf("Please enter a number between %d and %d: ", min, max);
    }
}

void show_progress(uint64_t current, uint64_t total, time_t start_time) {
    time_t now = time(NULL);
    double elapsed = difftime(now, start_time);
    
    if (elapsed < 1.0) elapsed = 1.0; // Avoid division by zero
    
    double percentage = (double)current / total * 100.0;
    double speed = (double)current / elapsed / (1024.0 * 1024.0); // MB/s
    
    uint64_t remaining_bytes = total - current;
    double eta = remaining_bytes / (current / elapsed);
    
    printf("\rProgress: %.1f%% (%.1f MB/s, ETA: %.0fs)    ", percentage, speed, eta);
    fflush(stdout);
}

int generate_pad(uint64_t size_bytes, int display_progress) {
    char temp_filename[32];
    char pad_filename[MAX_HASH_LENGTH + 10];
    char state_filename[MAX_HASH_LENGTH + 10];
    char hash_hex[MAX_HASH_LENGTH];
    
    // Create temporary filename
    snprintf(temp_filename, sizeof(temp_filename), "temp_%ld.pad", time(NULL));
    
    FILE* urandom = fopen("/dev/urandom", "rb");
    if (!urandom) {
        printf("Error: Cannot open /dev/urandom\n");
        return 1;
    }
    
    FILE* pad_file = fopen(temp_filename, "wb");
    if (!pad_file) {
        printf("Error: Cannot create temporary pad file %s\n", temp_filename);
        fclose(urandom);
        return 1;
    }
    
    unsigned char buffer[64 * 1024]; // 64KB buffer
    uint64_t bytes_written = 0;
    time_t start_time = time(NULL);
    
    if (display_progress) {
        printf("Generating pad...\n");
    }
    
    while (bytes_written < size_bytes) {
        uint64_t chunk_size = sizeof(buffer);
        if (size_bytes - bytes_written < chunk_size) {
            chunk_size = size_bytes - bytes_written;
        }
        
        if (fread(buffer, 1, (size_t)chunk_size, urandom) != (size_t)chunk_size) {
            printf("Error: Failed to read from /dev/urandom\n");
            fclose(urandom);
            fclose(pad_file);
            unlink(temp_filename);
            return 1;
        }
        
        if (fwrite(buffer, 1, (size_t)chunk_size, pad_file) != (size_t)chunk_size) {
            printf("Error: Failed to write to pad file\n");
            fclose(urandom);
            fclose(pad_file);
            unlink(temp_filename);
            return 1;
        }
        
        bytes_written += chunk_size;
        
        if (display_progress && bytes_written % PROGRESS_UPDATE_INTERVAL == 0) {
            show_progress(bytes_written, size_bytes, start_time);
        }
    }
    
    if (display_progress) {
        show_progress(size_bytes, size_bytes, start_time);
        printf("\n");
    }
    
    fclose(urandom);
    fclose(pad_file);
    
    // Calculate XOR checksum of the pad file
    if (calculate_checksum(temp_filename, hash_hex) != 0) {
        printf("Error: Cannot calculate pad checksum\n");
        unlink(temp_filename);
        return 1;
    }
    
    // Rename file to its hash
    snprintf(pad_filename, sizeof(pad_filename), "%s.pad", hash_hex);
    snprintf(state_filename, sizeof(state_filename), "%s.state", hash_hex);
    
    if (rename(temp_filename, pad_filename) != 0) {
        printf("Error: Cannot rename pad file to hash-based name\n");
        unlink(temp_filename);
        return 1;
    }
    
    // Set pad file to read-only
    if (chmod(pad_filename, S_IRUSR) != 0) {
        printf("Warning: Cannot set pad file to read-only\n");
    }
    
    // Initialize state file with offset 0
    if (write_state_offset(hash_hex, 0) != 0) {
        printf("Error: Failed to create state file\n");
        unlink(pad_filename);
        return 1;
    }
    
    double size_gb = (double)size_bytes / (1024.0 * 1024.0 * 1024.0);
    printf("Generated pad: %s (%.2f GB)\n", pad_filename, size_gb);
    printf("Pad hash: %s\n", hash_hex);
    printf("State file: %s\n", state_filename);
    printf("Pad file set to read-only\n");
    
    return 0;
}

int generate_pad_with_entropy(uint64_t size_bytes, int display_progress, int use_keyboard_entropy) {
    if (ensure_pads_directory() != 0) {
        printf("Error: Cannot create pads directory\n");
        return 1;
    }
    
    char temp_filename[64];
    char pad_path[MAX_HASH_LENGTH + 20];
    char state_path[MAX_HASH_LENGTH + 20];
    char hash_hex[MAX_HASH_LENGTH];
    
    // Create temporary filename
    snprintf(temp_filename, sizeof(temp_filename), "temp_%ld.pad", time(NULL));
    
    FILE* urandom = fopen("/dev/urandom", "rb");
    if (!urandom) {
        printf("Error: Cannot open /dev/urandom\n");
        return 1;
    }
    
    FILE* pad_file = fopen(temp_filename, "wb");
    if (!pad_file) {
        printf("Error: Cannot create temporary pad file %s\n", temp_filename);
        fclose(urandom);
        return 1;
    }
    
    // Setup keyboard entropy collection if requested
    struct termios original_termios;
    unsigned char* entropy_buffer = NULL;
    size_t entropy_collected = 0;
    int terminal_setup = 0;
    
    if (use_keyboard_entropy) {
        entropy_buffer = malloc(MAX_ENTROPY_BUFFER);
        if (!entropy_buffer) {
            printf("Error: Cannot allocate entropy buffer\n");
            fclose(urandom);
            fclose(pad_file);
            unlink(temp_filename);
            return 1;
        }
        
        if (setup_raw_terminal(&original_termios) == 0) {
            terminal_setup = 1;
            printf("Type random keys to add entropy (optional):\n");
        } else {
            printf("Warning: Cannot setup terminal for keyboard entropy collection\n");
            use_keyboard_entropy = 0;
            free(entropy_buffer);
            entropy_buffer = NULL;
        }
    }
    
    unsigned char urandom_buffer[64 * 1024]; // 64KB buffer
    unsigned char output_buffer[64 * 1024];
    uint64_t bytes_written = 0;
    
    if (display_progress) {
        printf("Generating pad...\n");
        if (use_keyboard_entropy) {
            printf("(Keyboard entropy: collecting...)\n");
        }
    }
    
    while (bytes_written < size_bytes) {
        uint64_t chunk_size = sizeof(urandom_buffer);
        if (size_bytes - bytes_written < chunk_size) {
            chunk_size = size_bytes - bytes_written;
        }
        
        // Read from /dev/urandom
        if (fread(urandom_buffer, 1, (size_t)chunk_size, urandom) != (size_t)chunk_size) {
            printf("Error: Failed to read from /dev/urandom\n");
            if (terminal_setup) restore_terminal(&original_termios);
            if (entropy_buffer) free(entropy_buffer);
            fclose(urandom);
            fclose(pad_file);
            unlink(temp_filename);
            return 1;
        }
        
        if (use_keyboard_entropy && terminal_setup) {
            // Collect available keyboard entropy
            size_t chunk_entropy = 0;
            collect_keyboard_entropy(entropy_buffer + entropy_collected, 
                                   MAX_ENTROPY_BUFFER - entropy_collected, &chunk_entropy);
            entropy_collected += chunk_entropy;
            
            if (entropy_collected > 512) { // Have enough entropy to mix
                // Copy urandom data to output buffer
                memcpy(output_buffer, urandom_buffer, chunk_size);
                
                // Simple XOR mixing with keyboard entropy
                simple_entropy_mix(output_buffer, chunk_size, entropy_buffer, entropy_collected);
                
                // Reset entropy buffer for next chunk
                entropy_collected = 0;
            } else {
                // Not enough entropy yet, use urandom only
                memcpy(output_buffer, urandom_buffer, chunk_size);
            }
        } else {
            // No keyboard entropy, use urandom directly
            memcpy(output_buffer, urandom_buffer, chunk_size);
        }
        
        if (fwrite(output_buffer, 1, (size_t)chunk_size, pad_file) != (size_t)chunk_size) {
            printf("Error: Failed to write to pad file\n");
            if (terminal_setup) restore_terminal(&original_termios);
            if (entropy_buffer) free(entropy_buffer);
            fclose(urandom);
            fclose(pad_file);
            unlink(temp_filename);
            return 1;
        }
        
        bytes_written += chunk_size;
        
        if (display_progress && bytes_written % PROGRESS_UPDATE_INTERVAL == 0) {
            printf("\rProgress: %.1f%% ", (double)bytes_written / size_bytes * 100.0);
            if (use_keyboard_entropy && terminal_setup) {
                printf("(keyboard entropy: %.1fKB) ", (double)entropy_collected / 1024.0);
            }
            fflush(stdout);
        }
    }
    
    if (terminal_setup) {
        restore_terminal(&original_termios);
    }
    if (entropy_buffer) {
        free(entropy_buffer);
    }
    
    if (display_progress) {
        printf("\rProgress: 100.0%%");
        if (use_keyboard_entropy) {
            printf(" (keyboard entropy: MIXED)");
        }
        printf("\n");
    }
    
    fclose(urandom);
    fclose(pad_file);
    
    // Calculate XOR checksum of the pad file
    if (calculate_checksum(temp_filename, hash_hex) != 0) {
        printf("Error: Cannot calculate pad checksum\n");
        unlink(temp_filename);
        return 1;
    }
    
    // Get final paths in pads directory
    get_pad_path(hash_hex, pad_path, state_path);
    
    if (rename(temp_filename, pad_path) != 0) {
        printf("Error: Cannot move pad file to pads directory\n");
        unlink(temp_filename);
        return 1;
    }
    
    // Set pad file to read-only
    if (chmod(pad_path, S_IRUSR) != 0) {
        printf("Warning: Cannot set pad file to read-only\n");
    }
    
    // Initialize state file with offset 0
    FILE* state_file = fopen(state_path, "wb");
    if (state_file) {
        uint64_t zero = 0;
        fwrite(&zero, sizeof(uint64_t), 1, state_file);
        fclose(state_file);
    } else {
        printf("Error: Failed to create state file\n");
        unlink(pad_path);
        return 1;
    }
    
    double size_gb = (double)size_bytes / (1024.0 * 1024.0 * 1024.0);
    printf("Generated pad: %s (%.2f GB)\n", pad_path, size_gb);
    printf("Pad checksum: %s\n", hash_hex);
    printf("State file: %s\n", state_path);
    if (use_keyboard_entropy) {
        printf("Enhanced with keyboard entropy!\n");
    }
    printf("Pad file set to read-only\n");
    
    return 0;
}

int encrypt_text(const char* pad_identifier) {
    char* pad_hash = find_pad_by_prefix(pad_identifier);
    if (!pad_hash) {
        return 1;
    }
    
    char input_text[MAX_INPUT_SIZE];
    char hash_hex[MAX_HASH_LENGTH];
    uint64_t current_offset;
    
    char pad_path[MAX_HASH_LENGTH + 20];
    char state_path[MAX_HASH_LENGTH + 20];
    get_pad_path(pad_hash, pad_path, state_path);
    
    // Check if pad file exists
    if (access(pad_path, R_OK) != 0) {
        printf("Error: Pad file %s not found\n", pad_path);
        free(pad_hash);
        return 1;
    }
    
    // Read current offset
    if (read_state_offset(pad_hash, &current_offset) != 0) {
        printf("Error: Cannot read state file\n");
        free(pad_hash);
        return 1;
    }
    
    // Calculate XOR checksum of pad file
    if (calculate_checksum(pad_path, hash_hex) != 0) {
        printf("Error: Cannot calculate pad checksum\n");
        free(pad_hash);
        return 1;
    }
    
    // Get input text from user
    printf("Enter text to encrypt: ");
    fflush(stdout);
    
    if (fgets(input_text, sizeof(input_text), stdin) == NULL) {
        printf("Error: Failed to read input\n");
        free(pad_hash);
        return 1;
    }
    
    // Remove newline if present
    size_t input_len = strlen(input_text);
    if (input_len > 0 && input_text[input_len - 1] == '\n') {
        input_text[input_len - 1] = '\0';
        input_len--;
    }
    
    if (input_len == 0) {
        printf("Error: No input provided\n");
        free(pad_hash);
        return 1;
    }
    
    // Check if we have enough pad space
    struct stat pad_stat;
    if (stat(pad_path, &pad_stat) != 0) {
        printf("Error: Cannot get pad file size\n");
        free(pad_hash);
        return 1;
    }
    
    if (current_offset + input_len > (uint64_t)pad_stat.st_size) {
        printf("Error: Not enough pad space remaining\n");
        printf("Need: %lu bytes, Available: %lu bytes\n", 
               input_len, (uint64_t)pad_stat.st_size - current_offset);
        free(pad_hash);
        return 1;
    }
    
    // Read pad data at current offset
    FILE* pad_file = fopen(pad_path, "rb");
    if (!pad_file) {
        printf("Error: Cannot open pad file\n");
        free(pad_hash);
        return 1;
    }
    
    if (fseek(pad_file, current_offset, SEEK_SET) != 0) {
        printf("Error: Cannot seek to offset in pad file\n");
        fclose(pad_file);
        free(pad_hash);
        return 1;
    }
    
    unsigned char* pad_data = malloc(input_len);
    if (fread(pad_data, 1, input_len, pad_file) != input_len) {
        printf("Error: Cannot read pad data\n");
        free(pad_data);
        fclose(pad_file);
        free(pad_hash);
        return 1;
    }
    fclose(pad_file);
    
    // XOR encrypt the input
    unsigned char* ciphertext = malloc(input_len);
    for (size_t i = 0; i < input_len; i++) {
        ciphertext[i] = input_text[i] ^ pad_data[i];
    }
    
    // Encode as base64
    char* base64_cipher = custom_base64_encode(ciphertext, input_len);
    
    // Update state offset
    if (write_state_offset(pad_hash, current_offset + input_len) != 0) {
        printf("Warning: Failed to update state file\n");
    }
    
    // Output in ASCII armor format
    printf("\n-----BEGIN OTP MESSAGE-----\n");
    printf("Version: %s\n", get_version());
    printf("Pad-Hash: %s\n", hash_hex);
    printf("Pad-Offset: %lu\n", current_offset);
    printf("\n");
    
    // Print base64 data in 64-character lines
    int b64_len = strlen(base64_cipher);
    for (int i = 0; i < b64_len; i += 64) {
        printf("%.64s\n", base64_cipher + i);
    }
    
    printf("-----END OTP MESSAGE-----\n\n");
    
    // Cleanup
    free(pad_data);
    free(ciphertext);
    free(base64_cipher);
    free(pad_hash);
    
    return 0;
}

int decrypt_text(const char* pad_identifier) {
    // For command line mode, pad_identifier is ignored - we'll get the hash from the message
    (void)pad_identifier; // Suppress unused parameter warning
    
    char line[MAX_LINE_LENGTH];
    char stored_hash[MAX_HASH_LENGTH];
    char current_hash[MAX_HASH_LENGTH];
    uint64_t pad_offset;
    char base64_data[MAX_INPUT_SIZE * 2] = {0};
    int in_data_section = 0;
    
    printf("Enter encrypted message (paste the full ASCII armor block):\n");
    
    // Read the ASCII armor format
    int found_begin = 0;
    while (fgets(line, sizeof(line), stdin)) {
        line[strcspn(line, "\n")] = 0;
        
        if (strcmp(line, "-----BEGIN OTP MESSAGE-----") == 0) {
            found_begin = 1;
            continue;
        }
        
        if (strcmp(line, "-----END OTP MESSAGE-----") == 0) {
            break;
        }
        
        if (!found_begin) continue;
        
        if (strncmp(line, "Pad-Hash: ", 10) == 0) {
            strncpy(stored_hash, line + 10, 64);
            stored_hash[64] = '\0';
        }
        else if (strncmp(line, "Pad-Offset: ", 12) == 0) {
            pad_offset = strtoull(line + 12, NULL, 10);
        }
        else if (strlen(line) == 0) {
            in_data_section = 1;
        }
        else if (in_data_section) {
            strncat(base64_data, line, sizeof(base64_data) - strlen(base64_data) - 1);
        }
    }
    
    if (!found_begin) {
        printf("Error: Invalid message format - missing BEGIN header\n");
        return 1;
    }
    
    // Now we have the pad hash from the message, construct filename
    char pad_path[MAX_HASH_LENGTH + 20];
    char state_path[MAX_HASH_LENGTH + 20];
    get_pad_path(stored_hash, pad_path, state_path);
    
    // Check if we have this pad
    if (access(pad_path, R_OK) != 0) {
        printf("Error: Required pad not found: %s\n", stored_hash);
        printf("Available pads:\n");
        list_available_pads();
        return 1;
    }
    
    // Verify pad integrity
    if (calculate_checksum(pad_path, current_hash) != 0) {
        printf("Error: Cannot calculate current pad checksum\n");
        return 1;
    }
    
    if (strcmp(stored_hash, current_hash) != 0) {
        printf("Warning: Pad integrity check failed!\n");
        printf("Expected: %s\n", stored_hash);
        printf("Current:  %s\n", current_hash);
        printf("Continue anyway? (y/N): ");
        fflush(stdout);
        
        char response[10];
        if (fgets(response, sizeof(response), stdin) == NULL || 
            (response[0] != 'y' && response[0] != 'Y')) {
            printf("Decryption aborted.\n");
            return 1;
        }
    } else {
        printf("Pad integrity: VERIFIED\n");
    }
    
    // Decode base64
    int ciphertext_len;
    unsigned char* ciphertext = custom_base64_decode(base64_data, &ciphertext_len);
    if (!ciphertext) {
        printf("Error: Invalid base64 data\n");
        return 1;
    }
    
    // Read pad data at specified offset
    FILE* pad_file = fopen(pad_path, "rb");
    if (!pad_file) {
        printf("Error: Cannot open pad file %s\n", pad_path);
        free(ciphertext);
        return 1;
    }
    
    if (fseek(pad_file, pad_offset, SEEK_SET) != 0) {
        printf("Error: Cannot seek to offset %lu in pad file\n", pad_offset);
        free(ciphertext);
        fclose(pad_file);
        return 1;
    }
    
    unsigned char* pad_data = malloc(ciphertext_len);
    if (fread(pad_data, 1, ciphertext_len, pad_file) != (size_t)ciphertext_len) {
        printf("Error: Cannot read pad data\n");
        free(ciphertext);
        free(pad_data);
        fclose(pad_file);
        return 1;
    }
    fclose(pad_file);
    
    // XOR decrypt
    char* plaintext = malloc(ciphertext_len + 1);
    for (int i = 0; i < ciphertext_len; i++) {
        plaintext[i] = ciphertext[i] ^ pad_data[i];
    }
    plaintext[ciphertext_len] = '\0';
    
    printf("Decrypted: %s\n", plaintext);
    
    // Cleanup
    free(ciphertext);
    free(pad_data);
    free(plaintext);
    
    return 0;
}

int read_state_offset(const char* pad_hash, uint64_t* offset) {
    char state_filename[MAX_HASH_LENGTH + 20];
    snprintf(state_filename, sizeof(state_filename), "%s/%s.state", PADS_DIR, pad_hash);
    
    FILE* state_file = fopen(state_filename, "rb");
    if (!state_file) {
        *offset = 0;
        return 0;
    }
    
    if (fread(offset, sizeof(uint64_t), 1, state_file) != 1) {
        fclose(state_file);
        *offset = 0;
        return 0;
    }
    
    fclose(state_file);
    return 0;
}

int write_state_offset(const char* pad_hash, uint64_t offset) {
    char state_filename[MAX_HASH_LENGTH + 20];
    snprintf(state_filename, sizeof(state_filename), "%s/%s.state", PADS_DIR, pad_hash);
    
    FILE* state_file = fopen(state_filename, "wb");
    if (!state_file) {
        return 1;
    }
    
    if (fwrite(&offset, sizeof(uint64_t), 1, state_file) != 1) {
        fclose(state_file);
        return 1;
    }
    
    fclose(state_file);
    return 0;
}

int calculate_checksum(const char* filename, char* checksum_hex) {
    FILE* file = fopen(filename, "rb");
    if (!file) {
        return 1;
    }
    
    unsigned char checksum[32];
    unsigned char buffer[64 * 1024];  // 64KB buffer for large files
    size_t bytes_read;
    
    // Initialize checksum
    memset(checksum, 0, 32);
    size_t total_bytes = 0;
    
    // Calculate XOR checksum of entire file
    while ((bytes_read = fread(buffer, 1, sizeof(buffer), file)) > 0) {
        // Process this chunk with XOR checksum
        for (size_t i = 0; i < bytes_read; i++) {
            unsigned char bucket = (total_bytes + i) % 32;
            checksum[bucket] ^= buffer[i] ^ (((total_bytes + i) >> 8) & 0xFF) ^ 
                               (((total_bytes + i) >> 16) & 0xFF) ^ (((total_bytes + i) >> 24) & 0xFF);
        }
        total_bytes += bytes_read;
    }
    
    fclose(file);
    
    // Now encrypt the checksum with the first 32 bytes of the pad
    fseek(file = fopen(filename, "rb"), 0, SEEK_SET);
    unsigned char pad_key[32];
    if (fread(pad_key, 1, 32, file) != 32) {
        fclose(file);
        return 1;
    }
    fclose(file);
    
    // XOR encrypt the checksum with pad data to create unique identifier
    unsigned char encrypted_checksum[32];
    for (int i = 0; i < 32; i++) {
        encrypted_checksum[i] = checksum[i] ^ pad_key[i];
    }
    
    // Convert to hex string (64 characters)
    for (int i = 0; i < 32; i++) {
        sprintf(checksum_hex + (i * 2), "%02x", encrypted_checksum[i]);
    }
    checksum_hex[64] = '\0';
    
    return 0;
}

// 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;
}

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);
}

int collect_keyboard_entropy(unsigned char* entropy_buffer, size_t max_size, size_t* collected) {
    struct timespec timestamp;
    unsigned char entropy_block[16];
    uint32_t sequence_counter = 0;
    char key;
    *collected = 0;
    
    while (*collected < max_size - 16) {
        if (read(STDIN_FILENO, &key, 1) == 1) {
            clock_gettime(CLOCK_MONOTONIC, &timestamp);
            
            // Create entropy block: [key][timestamp][sequence_counter]
            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, 3);
            
            // Add to entropy buffer
            memcpy(entropy_buffer + *collected, entropy_block, 16);
            *collected += 16;
            sequence_counter++;
        } else {
            // No key available, add some timing entropy
            clock_gettime(CLOCK_MONOTONIC, &timestamp);
            if (*collected + 12 < max_size) {
                memcpy(entropy_buffer + *collected, &timestamp, 12);
                *collected += 12;
            }
            usleep(1000); // 1ms delay
        }
    }
    
    return 0;
}


// Directory management functions
int ensure_pads_directory(void) {
    struct stat st = {0};
    if (stat(PADS_DIR, &st) == -1) {
        if (mkdir(PADS_DIR, 0755) != 0) {
            return 1;
        }
    }
    return 0;
}

void get_pad_path(const char* hash, char* pad_path, char* state_path) {
    snprintf(pad_path, MAX_HASH_LENGTH + 20, "%s/%s.pad", PADS_DIR, hash);
    snprintf(state_path, MAX_HASH_LENGTH + 20, "%s/%s.state", PADS_DIR, hash);
}


// Custom XOR checksum function
void xor_checksum_256(const unsigned char* data, size_t len, unsigned char checksum[32]) {
    memset(checksum, 0, 32);
    for (size_t i = 0; i < len; i++) {
        unsigned char bucket = i % 32;
        checksum[bucket] ^= data[i] ^ ((i >> 8) & 0xFF) ^ ((i >> 16) & 0xFF) ^ ((i >> 24) & 0xFF);
    }
}

// Custom base64 encode function
char* custom_base64_encode(const unsigned char* input, int length) {
    int output_length = 4 * ((length + 2) / 3);
    char* encoded = malloc(output_length + 1);
    if (!encoded) return NULL;
    
    int i, j;
    for (i = 0, j = 0; i < length;) {
        uint32_t octet_a = i < length ? input[i++] : 0;
        uint32_t octet_b = i < length ? input[i++] : 0;
        uint32_t octet_c = i < length ? input[i++] : 0;
        
        uint32_t triple = (octet_a << 16) + (octet_b << 8) + octet_c;
        
        encoded[j++] = base64_chars[(triple >> 18) & 63];
        encoded[j++] = base64_chars[(triple >> 12) & 63];
        encoded[j++] = base64_chars[(triple >> 6) & 63];
        encoded[j++] = base64_chars[triple & 63];
    }
    
    // Add padding
    for (int pad = 0; pad < (3 - length % 3) % 3; pad++) {
        encoded[output_length - 1 - pad] = '=';
    }
    
    encoded[output_length] = '\0';
    return encoded;
}

// Custom base64 decode function
unsigned char* custom_base64_decode(const char* input, int* output_length) {
    int input_length = strlen(input);
    if (input_length % 4 != 0) return NULL;
    
    *output_length = input_length / 4 * 3;
    if (input[input_length - 1] == '=') (*output_length)--;
    if (input[input_length - 2] == '=') (*output_length)--;
    
    unsigned char* decoded = malloc(*output_length);
    if (!decoded) return NULL;
    
    int i, j;
    for (i = 0, j = 0; i < input_length;) {
        int sextet_a = input[i] == '=' ? 0 & i++ : base64_decode_table[(unsigned char)input[i++]];
        int sextet_b = input[i] == '=' ? 0 & i++ : base64_decode_table[(unsigned char)input[i++]];
        int sextet_c = input[i] == '=' ? 0 & i++ : base64_decode_table[(unsigned char)input[i++]];
        int sextet_d = input[i] == '=' ? 0 & i++ : base64_decode_table[(unsigned char)input[i++]];
        
        if (sextet_a == -1 || sextet_b == -1 || sextet_c == -1 || sextet_d == -1) {
            free(decoded);
            return NULL;
        }
        
        uint32_t triple = (sextet_a << 18) + (sextet_b << 12) + (sextet_c << 6) + sextet_d;
        
        if (j < *output_length) decoded[j++] = (triple >> 16) & 255;
        if (j < *output_length) decoded[j++] = (triple >> 8) & 255;
        if (j < *output_length) decoded[j++] = triple & 255;
    }
    
    return decoded;
}

// Simple keyboard entropy mixing function
void simple_entropy_mix(unsigned char* urandom_buffer, size_t buffer_size, 
                        const unsigned char* entropy_data, size_t entropy_size) {
    if (!entropy_data || entropy_size == 0) return;
    
    for (size_t i = 0; i < buffer_size; i++) {
        // XOR with entropy data in a rotating pattern
        unsigned char entropy_byte = entropy_data[i % entropy_size];
        // Mix position information
        entropy_byte ^= (i & 0xFF) ^ ((i >> 8) & 0xFF);
        urandom_buffer[i] ^= entropy_byte;
    }
}

void print_usage(const char* program_name) {
    printf("OTP Cipher - One Time Pad Implementation %s\n", get_version());
    printf("%s\n", get_build_info());
    printf("Usage:\n");
    printf("  %s                              - Interactive mode\n", program_name);
    printf("  %s generate <size>              - Generate new pad\n", program_name);
    printf("  %s encrypt <pad_checksum_prefix> - Encrypt text\n", program_name);
    printf("  %s decrypt <pad_checksum_prefix> - Decrypt message\n", program_name);
    printf("  %s list                         - List available pads\n", program_name);
    printf("\nSize examples: 1GB, 5TB, 512MB, 2048 (bytes)\n");
    printf("Pad selection: Full hash, prefix, or number from list\n");
}