#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 "main.h"

// Initialize terminal dimensions
void init_terminal_dimensions(void) {
    struct winsize ws;

    // Try to get actual terminal size
    if (ioctl(STDOUT_FILENO, TIOCGWINSZ, &ws) == 0 && ws.ws_col > 0 && ws.ws_row > 0) {
        set_terminal_dimensions(ws.ws_col, ws.ws_row);
    }
    // If ioctl fails, keep the default values (80x24)
}

// Print centered header with = padding, screen clearing, and optional pause
void print_centered_header(const char* text, int pause_before_clear) {
    if (!text) return;

    // Phase 1: Pause if requested
    if (pause_before_clear) {
        printf("\nPress Enter to continue...");
        fflush(stdout);

        // Wait for Enter key
        int c;
        while ((c = getchar()) != '\n' && c != EOF) {
            // Consume any extra characters until newline
        }
    }

    // Phase 2: Clear screen using terminal height
    for (int i = 0; i < get_terminal_height(); i++) {
        printf("\n");
    }

    // Phase 3: Display centered header (existing logic)
    int text_len = strlen(text);
    int available_width = get_terminal_width();

    // Ensure minimum spacing: at least 1 space on each side
    int min_required = text_len + 4; // text + " " + text + " " (spaces around text)

    if (available_width < min_required) {
        // Terminal too narrow - just print the text with minimal formatting
        printf("=== %s ===\n", text);
        return;
    }

    // Calculate padding
    int total_padding = available_width - text_len - 2; // -2 for spaces around text
    int left_padding = total_padding / 2;
    int right_padding = total_padding - left_padding;

    // Print the header
    for (int i = 0; i < left_padding; i++) {
        printf("=");
    }
    printf(" %s ", text);
    for (int i = 0; i < right_padding; i++) {
        printf("=");
    }
    printf("\n");
}

// Interactive mode main loop
int interactive_mode(void) {
    char input[10];
        printf("\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n");
    while (1) {
        show_main_menu();

        if (!fgets(input, sizeof(input), stdin)) {
            printf("Goodbye!\n");
            break;
        }

        char choice = toupper(input[0]);

        switch (choice) {
            case 'T':
                handle_text_encrypt();
                break;
            case 'F':
                handle_file_encrypt();
                break;
            case 'R':
                handle_directory_encrypt();
                break;
            case 'D':
                handle_decrypt_menu();
                break;
            case 'P':
                handle_pads_menu();
                break;
            case 'X':
            case 'Q':
                printf("Goodbye!\n");
                return 0;
            default:
                printf("Invalid choice. Please try again.\n");
                break;
        }
    }

    return 0;
}

void show_main_menu(void) {
    printf("\n");
    char header[64];
    snprintf(header, sizeof(header), "Main Menu - OTP %s", OTP_VERSION);
    print_centered_header(header, 0);
    printf("\n");

    printf(" \033[4mT\033[0mext encrypt\n");       //TEXT ENCRYPT
    printf(" \033[4mF\033[0mile encrypt\n");       //FILE ENCRYPT
    printf(" Di\033[4mr\033[0mectory encrypt\n");  //DIRECTORY ENCRYPT
    printf(" \033[4mD\033[0mecrypt\n");            //DECRYPT
    printf(" \033[4mP\033[0mads\n");               //PADS
    printf(" E\033[4mx\033[0mit\n");               //EXIT
    printf("\nSelect option: ");
}

int handle_generate_menu(void) {
    printf("\n");
    print_centered_header("Generate New Pad", 0);
    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;
    }

    double size_gb = (double)size / (1000.0 * 1000.0 * 1000.0);
    printf("Generating %.2f GB pad...\n", size_gb);
    printf("Note: Use 'Add entropy' in Pads menu to enhance randomness after creation.\n");

    return generate_pad(size, 1);
}

int handle_encrypt_menu(void) {
    printf("\n");
    print_centered_header("Encrypt Data", 0);

    printf("Available pads:\n");
    char* selected = select_pad_interactive("Available pads:", "Select pad (or press Enter to continue)", PAD_FILTER_ALL, 0);
    int pad_count = 1; // Assume at least 1 pad if function returned
    if (selected) {
        free(selected);
    } else {
        pad_count = 0;
    }
    if (pad_count == 0) {
        printf("No pads available. Generate a pad first.\n");
        return 1;
    }

    // Ask user to choose between text and file encryption
    printf("\nSelect encryption type:\n");
    printf(" 1. Text message\n");
    printf(" 2. File\n");
    printf("Enter choice (1-2) or 'esc' to cancel: ");

    char choice_input[10];
    if (!fgets(choice_input, sizeof(choice_input), stdin)) {
        printf("Error: Failed to read input\n");
        return 1;
    }

    choice_input[strcspn(choice_input, "\n")] = 0;
    
    // Check for ESC/cancel
    if (is_escape_input(choice_input)) {
        printf("Returning to main menu...\n");
        return 0;
    }

    int choice = atoi(choice_input);

    if (choice == 1) {
        // Text encryption - use unified pad selection
        char* selected_pad = select_pad_interactive("Select Pad for Text Encryption",
                                                  "Select pad (by prefix)",
                                                  PAD_FILTER_ALL, 1);
        if (!selected_pad) {
            printf("Text encryption cancelled.\n");
            return 1;
        }

        int result = encrypt_text(selected_pad, NULL);  // NULL for interactive mode
        free(selected_pad);
        return result;
    }
    else if (choice == 2) {
        // File encryption
        printf("\nFile selection options:\n");
        printf(" 1. Type file path directly\n");
        printf(" 2. Use file manager\n");
        printf("Enter choice (1-2) or 'esc' to cancel: ");

        char file_choice[10];
        char input_file[512];

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

        file_choice[strcspn(file_choice, "\n")] = 0;
        
        // Check for ESC/cancel
        if (is_escape_input(file_choice)) {
            printf("Returning to main menu...\n");
            return 0;
        }

        if (atoi(file_choice) == 2) {
            // Use file manager
            if (launch_file_manager(".", input_file, sizeof(input_file)) != 0) {
                printf("Falling back to manual file path entry.\n");
                printf("Enter input file path (or 'esc' to cancel): ");
                if (!fgets(input_file, sizeof(input_file), stdin)) {
                    printf("Error: Failed to read input\n");
                    return 1;
                }
                input_file[strcspn(input_file, "\n")] = 0;
                
                // Check for ESC/cancel
                if (is_escape_input(input_file)) {
                    printf("Returning to main menu...\n");
                    return 0;
                }
            }
        } else {
            // Direct file path input
            printf("Enter input file path (or 'esc' to cancel): ");
            if (!fgets(input_file, sizeof(input_file), stdin)) {
                printf("Error: Failed to read input\n");
                return 1;
            }
            input_file[strcspn(input_file, "\n")] = 0;
            
            // Check for ESC/cancel
            if (is_escape_input(input_file)) {
                printf("Returning to main menu...\n");
                return 0;
            }
        }

        // Check if file exists
        if (access(input_file, R_OK) != 0) {
            printf("Error: File '%s' not found or cannot be read\n", input_file);
            return 1;
        }

        // Use unified pad selection
        char* selected_pad = select_pad_interactive("Select Pad for File Encryption",
                                                  "Select pad (by prefix)",
                                                  PAD_FILTER_ALL, 1);
        if (!selected_pad) {
            printf("File encryption cancelled.\n");
            return 1;
        }

        // Ask for output format
        printf("\nSelect output format:\n");
        printf(" 1. Binary (.otp) - preserves file permissions\n");
        printf(" 2. ASCII (.otp.asc) - text-safe format\n");
        printf("Enter choice (1-2) or 'esc' to cancel: ");

        char format_input[10];
        if (!fgets(format_input, sizeof(format_input), stdin)) {
            printf("Error: Failed to read input\n");
            free(selected_pad);
            return 1;
        }

        format_input[strcspn(format_input, "\n")] = 0;
        
        // Check for ESC/cancel
        if (is_escape_input(format_input)) {
            printf("Returning to main menu...\n");
            free(selected_pad);
            return 0;
        }

        int ascii_armor = (atoi(format_input) == 2) ? 1 : 0;

        // Generate default output filename with files directory and use enhanced input function
        char default_output[1024]; // Increased size to prevent truncation warnings
        char temp_default[1024];

        // Generate base filename with appropriate extension
        if (ascii_armor) {
            snprintf(temp_default, sizeof(temp_default), "%s.otp.asc", input_file);
        } else {
            snprintf(temp_default, sizeof(temp_default), "%s.otp", input_file);
        }

        // Apply files directory default path
        get_default_file_path(temp_default, default_output, sizeof(default_output));

        char output_file[512];
        if (get_filename_with_default("Output filename:", default_output, output_file, sizeof(output_file)) != 0) {
            printf("Error: Failed to read input\n");
            return 1;
        }

        const char* output_filename = output_file;

        int result = encrypt_file(selected_pad, input_file, output_filename, ascii_armor);
        free(selected_pad);
        return result;
    }
    else {
        printf("Invalid choice. Please enter 1 or 2.\n");
        return 1;
    }
}

int handle_decrypt_menu(void) {
    printf("\n");
    print_centered_header("Smart Decrypt", 0);
    printf("Enter encrypted data (paste ASCII armor), file path, or press Enter to browse files:\n");
    printf("(Type 'esc' or 'q' to return to main menu)\n");

    char input_line[MAX_LINE_LENGTH];
    if (!fgets(input_line, sizeof(input_line), stdin)) {
        printf("Error: Failed to read input\n");
        return 1;
    }

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

    // Check for ESC/cancel
    if (is_escape_input(input_line)) {
        printf("Returning to main menu...\n");
        return 0;
    }

    // Trim leading whitespace to handle pasted content better
    char* trimmed_input = input_line;
    while (*trimmed_input == ' ' || *trimmed_input == '\t') {
        trimmed_input++;
    }

    // Check for ASCII armor FIRST, before checking for empty input
    // This handles cases where pasted text starts with the header
    if (strncmp(trimmed_input, "-----BEGIN OTP MESSAGE-----", 27) == 0) {
        // Looks like ASCII armor - collect the full message
        char full_message[MAX_INPUT_SIZE * 4] = {0};
        strcat(full_message, input_line);
        strcat(full_message, "\n");

        printf("Continue pasting the message (end with -----END OTP MESSAGE-----):\n");

        char line[MAX_LINE_LENGTH];
        while (fgets(line, sizeof(line), stdin)) {
            strncat(full_message, line, sizeof(full_message) - strlen(full_message) - 1);
            if (strstr(line, "-----END OTP MESSAGE-----")) {
                break;
            }
        }

        return decrypt_text(NULL, full_message);
    }
    else if (strlen(trimmed_input) == 0) {
        // Empty input - launch file manager to browse for files
        char selected_file[512];
        if (launch_file_manager(get_files_directory(), selected_file, sizeof(selected_file)) != 0) {
            printf("File browsing cancelled or failed.\n");
            printf("Returning to main menu...\n");
            return 0;
        }

        // Generate smart default output filename with files directory and use enhanced input function
        char temp_default[512];
        char default_output[512];
        strncpy(temp_default, selected_file, sizeof(temp_default) - 1);
        temp_default[sizeof(temp_default) - 1] = '\0';

        // Remove common encrypted extensions to get a better default
        if (strstr(temp_default, ".tar.gz.otp")) {
            // Directory archive - remove .tar.gz.otp to get original directory name
            char* ext_pos = strstr(temp_default, ".tar.gz.otp");
            *ext_pos = '\0';
        } else if (strstr(temp_default, ".tar.otp")) {
            // Directory archive without compression - remove .tar.otp
            char* ext_pos = strstr(temp_default, ".tar.otp");
            *ext_pos = '\0';
        } else if (strstr(temp_default, ".otp.asc")) {
            // Replace .otp.asc with original extension or no extension
            char* ext_pos = strstr(temp_default, ".otp.asc");
            *ext_pos = '\0';
        } else if (strstr(temp_default, ".otp")) {
            // Replace .otp with original extension or no extension
            char* ext_pos = strstr(temp_default, ".otp");
            *ext_pos = '\0';
        } else {
            // No recognized encrypted extension, add .decrypted suffix
            strncat(temp_default, ".decrypted", sizeof(temp_default) - strlen(temp_default) - 1);
        }

        // Apply files directory default path
        get_default_file_path(temp_default, default_output, sizeof(default_output));

        char output_file[512];
        if (get_filename_with_default("Output filename:", default_output, output_file, sizeof(output_file)) != 0) {
            printf("Error: Failed to read input\n");
            return 1;
        }

        // Check if it's a directory archive
        if (strstr(selected_file, ".tar.gz.otp") || strstr(selected_file, ".tar.otp")) {
            // It's a directory archive - extract to directory
            char extract_dir[512];
            strncpy(extract_dir, output_file, sizeof(extract_dir) - 1);
            extract_dir[sizeof(extract_dir) - 1] = '\0';
            
            // Remove .tar.gz.otp or .tar.otp extension to get directory name
            char* ext = strstr(extract_dir, ".tar.gz.otp");
            if (!ext) ext = strstr(extract_dir, ".tar.otp");
            if (ext) *ext = '\0';
            
            printf("Extracting directory archive to: %s/\n", extract_dir);
            return decrypt_and_extract_directory(selected_file, extract_dir);
        } else {
            return decrypt_file(selected_file, output_file);
        }
    }
    else {
        // Check if it looks like a file path
        if (access(trimmed_input, R_OK) == 0) {
            // It's a valid file - decrypt it with enhanced input for output filename
            char temp_default[512];
            char default_output[512];
            strncpy(temp_default, trimmed_input, sizeof(temp_default) - 1);
            temp_default[sizeof(temp_default) - 1] = '\0';

            // Remove common encrypted extensions to get a better default
            if (strstr(temp_default, ".tar.gz.otp")) {
                // Directory archive - remove .tar.gz.otp to get original directory name
                char* ext_pos = strstr(temp_default, ".tar.gz.otp");
                *ext_pos = '\0';
            } else if (strstr(temp_default, ".tar.otp")) {
                // Directory archive without compression - remove .tar.otp
                char* ext_pos = strstr(temp_default, ".tar.otp");
                *ext_pos = '\0';
            } else if (strstr(temp_default, ".otp.asc")) {
                // Replace .otp.asc with original extension or no extension
                char* ext_pos = strstr(temp_default, ".otp.asc");
                *ext_pos = '\0';
            } else if (strstr(temp_default, ".otp")) {
                // Replace .otp with original extension or no extension
                char* ext_pos = strstr(temp_default, ".otp");
                *ext_pos = '\0';
            } else {
                // No recognized encrypted extension, add .decrypted suffix
                strncat(temp_default, ".decrypted", sizeof(temp_default) - strlen(temp_default) - 1);
            }

            // Apply files directory default path
            get_default_file_path(temp_default, default_output, sizeof(default_output));

            char output_file[512];
            if (get_filename_with_default("Output filename:", default_output, output_file, sizeof(output_file)) != 0) {
                printf("Error: Failed to read input\n");
                return 1;
            }

            // Check if it's a directory archive
            if (strstr(trimmed_input, ".tar.gz.otp") || strstr(trimmed_input, ".tar.otp")) {
                // It's a directory archive - extract to directory
                char extract_dir[512];
                strncpy(extract_dir, output_file, sizeof(extract_dir) - 1);
                extract_dir[sizeof(extract_dir) - 1] = '\0';
                
                // Remove .tar.gz.otp or .tar.otp extension to get directory name
                char* ext = strstr(extract_dir, ".tar.gz.otp");
                if (!ext) ext = strstr(extract_dir, ".tar.otp");
                if (ext) *ext = '\0';
                
                printf("Extracting directory archive to: %s/\n", extract_dir);
                return decrypt_and_extract_directory(trimmed_input, extract_dir);
            } else {
                return decrypt_file(trimmed_input, output_file);
            }
        } else {
            printf("Input not recognized as ASCII armor or valid file path.\n");
            return 1;
        }
    }
}

int handle_text_encrypt(void) {
    printf("\n");
    print_centered_header("Text Encrypt", 0);

    // Launch text editor directly
    char text_buffer[MAX_INPUT_SIZE];
    if (launch_text_editor(NULL, text_buffer, sizeof(text_buffer)) != 0) {
        printf("Error: Could not launch text editor\n");
        return 1;
    }

    if (strlen(text_buffer) == 0) {
        printf("No text entered - canceling encryption\n");
        return 1;
    }

    // Use unified pad selection
    char* selected_pad = select_pad_interactive("Select Pad for Text Encryption",
                                              "Select pad (by prefix)",
                                              PAD_FILTER_ALL, 1);
    if (!selected_pad) {
        printf("Text encryption cancelled.\n");
        return 1;
    }

    int result = encrypt_text(selected_pad, text_buffer);
    free(selected_pad);
    return result;
}

int handle_file_encrypt(void) {
    printf("\n");
    print_centered_header("File Encrypt", 0);

    // Launch file manager directly
    char input_file[512];
    if (launch_file_manager(".", input_file, sizeof(input_file)) != 0) {
        printf("Error: Could not launch file manager\n");
        return 1;
    }

    // Check if file exists
    if (access(input_file, R_OK) != 0) {
        printf("Error: File '%s' not found or cannot be read\n", input_file);
        return 1;
    }

    // Use unified pad selection
    char* selected_pad = select_pad_interactive("Select Pad for File Encryption",
                                              "Select pad (by prefix)",
                                              PAD_FILTER_ALL, 1);
    if (!selected_pad) {
        printf("File encryption cancelled.\n");
        return 1;
    }

    // Ask for output format
    printf("\nSelect output format:\n");
    printf(" 1. Binary (.otp) - preserves file permissions\n");
    printf(" 2. ASCII (.otp.asc) - text-safe format\n");
    printf("Enter choice (1-2) or 'esc' to cancel: ");

    char format_input[10];
    if (!fgets(format_input, sizeof(format_input), stdin)) {
        printf("Error: Failed to read input\n");
        free(selected_pad);
        return 1;
    }

    format_input[strcspn(format_input, "\n")] = 0;
    
    // Check for ESC/cancel
    if (is_escape_input(format_input)) {
        printf("Returning to main menu...\n");
        free(selected_pad);
        return 0;
    }

    int ascii_armor = (atoi(format_input) == 2) ? 1 : 0;

    // Generate default output filename
    char default_output[1024]; // Increased buffer size to prevent truncation warnings
    if (ascii_armor) {
        snprintf(default_output, sizeof(default_output), "%s.otp.asc", input_file);
    } else {
        snprintf(default_output, sizeof(default_output), "%s.otp", input_file);
    }

    // Use enhanced input function for output filename
    char output_file[512];
    if (get_filename_with_default("Output filename:", default_output, output_file, sizeof(output_file)) != 0) {
        printf("Error: Failed to read input\n");
        return 1;
    }

    const char* output_filename = output_file;

    int result = encrypt_file(selected_pad, input_file, output_filename, ascii_armor);
    free(selected_pad);
    return result;
}

// Comparison function for qsort (case-insensitive)
static int compare_dir_strings(const void *a, const void *b) {
    return strcasecmp(*(const char**)a, *(const char**)b);
}

// Function to get list of subdirectories
static int get_subdirs(const char *path, char ***subdirs) {
    DIR *dir = opendir(path);
    if (!dir) return 0;
    
    int count = 0;
    int capacity = 10;
    *subdirs = malloc(capacity * sizeof(char*));
    
    struct dirent *entry;
    while ((entry = readdir(dir)) != NULL) {
        if (entry->d_type == DT_DIR) {
            // Skip . and ..
            if (strcmp(entry->d_name, ".") == 0 || strcmp(entry->d_name, "..") == 0)
                continue;
            
            if (count >= capacity) {
                capacity *= 2;
                *subdirs = realloc(*subdirs, capacity * sizeof(char*));
            }
            (*subdirs)[count++] = strdup(entry->d_name);
        }
    }
    closedir(dir);
    
    // Sort alphabetically
    if (count > 0) {
        qsort(*subdirs, count, sizeof(char*), compare_dir_strings);
    }
    
    return count;
}

// Function to get a single keypress without echo
static int getch_nav(void) {
    struct termios oldt, newt;
    int ch;
    tcgetattr(STDIN_FILENO, &oldt);
    newt = oldt;
    newt.c_lflag &= ~(ICANON | ECHO);
    tcsetattr(STDIN_FILENO, TCSANOW, &newt);
    ch = getchar();
    tcsetattr(STDIN_FILENO, TCSANOW, &oldt);
    return ch;
}

// Navigate directories with arrow keys
static char* navigate_directory_interactive(void) {
    char current_path[PATH_MAX];
    char original_path[PATH_MAX];
    getcwd(original_path, sizeof(original_path));
    getcwd(current_path, sizeof(current_path));
    
    char **subdirs = NULL;
    int subdir_count = 0;
    int current_index = 0;
    
    printf("\nNavigate with arrow keys (UP/DOWN: cycle, RIGHT: enter, LEFT: back, ENTER: select, Q: cancel)\n");
    printf("\033[?25l"); // Hide cursor
    
    while (1) {
        // Clear line and show current path
        printf("\r\033[K%s", current_path);
        
        // If we have subdirectories and an index, show current selection
        if (subdir_count > 0 && current_index < subdir_count) {
            printf("/%s", subdirs[current_index]);
        }
        
        fflush(stdout);
        
        int ch = getch_nav();
        
        // Handle arrow keys (they come as escape sequences)
        if (ch == 27) { // ESC sequence
            getch_nav(); // Skip '['
            ch = getch_nav();
            
            if (ch == 'A') { // UP arrow
                // Load subdirs if not loaded
                if (subdirs == NULL) {
                    subdir_count = get_subdirs(current_path, &subdirs);
                    current_index = 0;
                } else if (subdir_count > 0) {
                    current_index = (current_index - 1 + subdir_count) % subdir_count;
                }
            }
            else if (ch == 'B') { // DOWN arrow
                // Load subdirs if not loaded
                if (subdirs == NULL) {
                    subdir_count = get_subdirs(current_path, &subdirs);
                    current_index = 0;
                } else if (subdir_count > 0) {
                    current_index = (current_index + 1) % subdir_count;
                }
            }
            else if (ch == 'C') { // RIGHT arrow - go deeper
                if (subdir_count > 0 && current_index < subdir_count) {
                    // Navigate into selected directory
                    char new_path[PATH_MAX];
                    snprintf(new_path, sizeof(new_path), "%s/%s", current_path, subdirs[current_index]);
                    
                    if (chdir(new_path) == 0) {
                        getcwd(current_path, sizeof(current_path));
                        
                        // Free old subdirs
                        for (int i = 0; i < subdir_count; i++) {
                            free(subdirs[i]);
                        }
                        free(subdirs);
                        subdirs = NULL;
                        
                        // Load subdirs of new directory and show first one
                        subdir_count = get_subdirs(current_path, &subdirs);
                        current_index = 0;
                    }
                }
            }
            else if (ch == 'D') { // LEFT arrow - go up
                if (chdir("..") == 0) {
                    getcwd(current_path, sizeof(current_path));
                    
                    // Free old subdirs
                    for (int i = 0; i < subdir_count; i++) {
                        free(subdirs[i]);
                    }
                    free(subdirs);
                    subdirs = NULL;
                    subdir_count = 0;
                    current_index = 0;
                }
            }
        }
        else if (ch == '\n' || ch == '\r') { // ENTER - confirm selection
            // If a subdirectory is displayed, navigate into it first
            if (subdir_count > 0 && current_index < subdir_count) {
                char new_path[PATH_MAX];
                snprintf(new_path, sizeof(new_path), "%s/%s", current_path, subdirs[current_index]);
                
                if (chdir(new_path) == 0) {
                    getcwd(current_path, sizeof(current_path));
                }
            }
            break;
        }
        else if (ch == 'q' || ch == 'Q') { // Q to quit
            printf("\033[?25h\n"); // Show cursor
            // Restore original directory
            chdir(original_path);
            // Clean up
            for (int i = 0; i < subdir_count; i++) {
                free(subdirs[i]);
            }
            free(subdirs);
            return NULL;
        }
    }
    
    printf("\033[?25h\n"); // Show cursor
    
    // Clean up
    for (int i = 0; i < subdir_count; i++) {
        free(subdirs[i]);
    }
    free(subdirs);
    
    // Restore original directory before returning
    char* result = strdup(current_path);
    chdir(original_path);
    
    return result;
}

int handle_directory_encrypt(void) {
    printf("\n");
    print_centered_header("Directory Encrypt", 0);

    // Directory selection options
    printf("\nDirectory selection options:\n");
    printf(" 1. Type directory path directly\n");
    printf(" 2. Navigate with arrow keys\n");
    printf(" 3. Use file manager (navigate to directory)\n");
    printf("Enter choice (1-3) or 'esc' to cancel: ");

    char choice_input[10];
    char dir_path[512];

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

    choice_input[strcspn(choice_input, "\n")] = 0;
    
    // Check for ESC/cancel
    if (is_escape_input(choice_input)) {
        printf("Returning to main menu...\n");
        return 0;
    }

    int choice = atoi(choice_input);
    
    if (choice == 2) {
        // Use arrow key navigation
        char* selected = navigate_directory_interactive();
        if (!selected) {
            printf("Directory selection cancelled.\n");
            return 0;
        }
        strncpy(dir_path, selected, sizeof(dir_path) - 1);
        dir_path[sizeof(dir_path) - 1] = '\0';
        free(selected);
        printf("Selected: %s\n", dir_path);
    }
    else if (choice == 3) {
        // Use directory manager
        if (launch_directory_manager(".", dir_path, sizeof(dir_path)) != 0) {
            printf("Falling back to manual directory path entry.\n");
            printf("Enter directory path to encrypt (or 'esc' to cancel): ");
            if (!fgets(dir_path, sizeof(dir_path), stdin)) {
                printf("Error: Failed to read input\n");
                return 1;
            }
            dir_path[strcspn(dir_path, "\n")] = 0;
            
            // Check for ESC/cancel
            if (is_escape_input(dir_path)) {
                printf("Returning to main menu...\n");
                return 0;
            }
        }
    } else {
        // Direct directory path input
        printf("Enter directory path to encrypt (or 'esc' to cancel): ");
        if (!fgets(dir_path, sizeof(dir_path), stdin)) {
            printf("Error: Failed to read input\n");
            return 1;
        }
        dir_path[strcspn(dir_path, "\n")] = 0;
        
        // Check for ESC/cancel
        if (is_escape_input(dir_path)) {
            printf("Returning to main menu...\n");
            return 0;
        }
    }

    // Check if directory exists
    struct stat st;
    if (stat(dir_path, &st) != 0 || !S_ISDIR(st.st_mode)) {
        printf("Error: '%s' is not a valid directory\n", dir_path);
        return 1;
    }

    // Select pad
    char* selected_pad = select_pad_interactive("Select Pad for Directory Encryption",
                                              "Select pad (by prefix)",
                                              PAD_FILTER_ALL, 1);
    if (!selected_pad) {
        printf("Directory encryption cancelled.\n");
        return 1;
    }

    // Generate default output filename - append .tar.gz.otp to the directory path
    char default_output[1024];
    
    // Remove trailing slash if present
    char clean_path[512];
    strncpy(clean_path, dir_path, sizeof(clean_path) - 1);
    clean_path[sizeof(clean_path) - 1] = '\0';
    
    size_t path_len = strlen(clean_path);
    if (path_len > 0 && clean_path[path_len - 1] == '/') {
        clean_path[path_len - 1] = '\0';
    }
    
    snprintf(default_output, sizeof(default_output), "%s.tar.gz.otp", clean_path);

    // Get output filename
    char output_file[512];
    if (get_filename_with_default("Output filename:", default_output, output_file, sizeof(output_file)) != 0) {
        printf("Error: Failed to read input\n");
        free(selected_pad);
        return 1;
    }

    // Encrypt directory
    int result = encrypt_directory(dir_path, selected_pad, output_file);
    free(selected_pad);
    
    return result;
}