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

// Basic TrueRNG entropy collection function
int collect_truerng_entropy(unsigned char* entropy_buffer, size_t target_bytes, size_t* collected_bytes, int display_progress) {
    hardware_rng_device_t devices[10];
    int num_devices_found = 0;

    // Detect available TrueRNG devices
    if (detect_all_hardware_rng_devices(devices, 10, &num_devices_found) != 0) {
        if (display_progress) {
            printf("Error: Failed to detect hardware RNG devices\n");
        }
        return 1;
    }

    if (num_devices_found == 0) {
        if (display_progress) {
            printf("No hardware RNG devices found.\n");
            printf("\nSupported devices:\n");
            printf("  - TrueRNG Original (VID: %s, PID: %s)\n", TRUERNG_VID, TRUERNG_ORIGINAL_PID);
            printf("  - TrueRNG Pro (VID: %s, PID: %s)\n", TRUERNG_VID, TRUERNG_PRO_PID);
            printf("  - TrueRNG Pro V2 (VID: %s, PID: %s)\n", TRUERNG_VID, TRUERNG_PRO_V2_PID);
            printf("\nPlease connect a TrueRNG device and try again.\n");
        }
        return 1;
    }

    // Use first available device
    hardware_rng_device_t* selected_device = &devices[0];

    if (display_progress) {
        printf("Using device: %s\n", selected_device->friendly_name);
        printf("Collecting %zu bytes of entropy...\n", target_bytes);
    }

    // Collect entropy from the device
    int result = collect_truerng_entropy_from_device(selected_device, entropy_buffer, target_bytes, collected_bytes, display_progress);

    if (result != 0) {
        if (display_progress) {
            printf("Error: Failed to collect entropy from TrueRNG device\n");
        }
        return 1;
    }

    if (display_progress) {
        printf("✓ Successfully collected %zu bytes of entropy from TrueRNG device\n", *collected_bytes);
    }

    return 0;
}

// Streaming entropy collection directly to pad file
int collect_truerng_entropy_streaming_from_device(const hardware_rng_device_t* device, const char* pad_chksum,
                                                  size_t total_bytes, int display_progress, int entropy_mode) {
    (void)entropy_mode; // Suppress unused parameter warning
    if (!device || !pad_chksum || total_bytes == 0) {
        return 1; // Invalid parameters
    }


    // 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) {
        if (display_progress) {
            printf("Error: Pad file not found: %s\n", pad_path);
        }
        return 1;
    }

    uint64_t pad_size = pad_stat.st_size;
    if (total_bytes > pad_size) {
        if (display_progress) {
            printf("Error: Requested entropy (%zu bytes) exceeds pad size (%lu bytes)\n", total_bytes, pad_size);
        }
        return 1;
    }

    // Open the RNG device
    int device_fd = open(device->port_path, O_RDONLY | O_NOCTTY);
    if (device_fd < 0) {
        if (display_progress) {
            printf("Error: Cannot open RNG device %s: %s\n", device->port_path, strerror(errno));
        }
        return 1;
    }

    // Configure serial port for this device type
    if (configure_rng_serial_port(device_fd, device->device_type) != 0) {
        if (display_progress) {
            printf("Error: Failed to configure serial port for %s\n", device->friendly_name);
        }
        close(device_fd);
        return 1;
    }

    // Standard delay for TrueRNG devices
    usleep(100000); // 100ms

    // Open pad file for read/write
    FILE* pad_file = fopen(pad_path, "r+b");
    if (!pad_file) {
        if (display_progress) {
            printf("Error: Cannot open pad file for modification: %s\n", pad_path);
            printf("Reason: %s\n", strerror(errno));

            // Provide additional diagnostics
            if (errno == EROFS) {
                printf("The filesystem appears to be read-only. Check if the drive is mounted read-only.\n");
            } else if (errno == EACCES) {
                printf("Permission denied. Check file permissions and mount options.\n");
            } else if (errno == ENOENT) {
                printf("File not found. The pad file may have been moved or deleted.\n");
            } else if (errno == EISDIR) {
                printf("Path is a directory, not a file.\n");
            } else {
                printf("This may be due to filesystem limitations or mount options.\n");
            }

            printf("\nTroubleshooting suggestions:\n");
            printf("1. Ensure the external drive is mounted read-write: mount -o remount,rw /media/teknari/OTP_01\n");
            printf("2. Check file permissions: ls -la '%s'\n", pad_path);
            printf("3. Verify the drive supports the required operations\n");
            printf("4. Try copying the pad to local storage, enhancing it, then copying back\n");
        }
        close(device_fd);
        return 1;
    }

    if (display_progress) {
        printf("Streaming entropy from %s to pad...\n", device->friendly_name);
        printf("Pad size: %.2f GB (%lu bytes)\n", (double)pad_size / (1000.0*1000.0*1000.0), pad_size);
        printf("Enhancing entire pad with hardware entropy\n");
    }

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

    while (bytes_processed < total_bytes && !error_occurred) {
        size_t chunk_size = sizeof(buffer);
        if (total_bytes - bytes_processed < chunk_size) {
            chunk_size = total_bytes - bytes_processed;
        }

        // Read entropy from device
        ssize_t entropy_read = read(device_fd, buffer, chunk_size);
        if (entropy_read < 0) {
            if (errno == EINTR) {
                continue; // Interrupted, try again
            }
            if (display_progress) {
                printf("Error: Failed to read from TrueRNG device: %s\n", strerror(errno));
                printf("Device may have been disconnected during operation.\n");
            }
            error_occurred = 1;
            break;
        }

        if (entropy_read == 0) {
            if (display_progress) {
                printf("Error: TrueRNG device returned no data (device disconnected?)\n");
            }
            error_occurred = 1;
            break;
        }

        // Read current pad data at this position
        if (fseek(pad_file, bytes_processed, SEEK_SET) != 0) {
            if (display_progress) {
                printf("Error: Cannot seek to position %zu in pad file\n", bytes_processed);
            }
            error_occurred = 1;
            break;
        }

        unsigned char pad_data[64 * 1024];
        size_t pad_read = fread(pad_data, 1, entropy_read, pad_file);
        if (pad_read != (size_t)entropy_read) {
            if (display_progress) {
                printf("Error: Cannot read pad data at position %zu\n", bytes_processed);
            }
            error_occurred = 1;
            break;
        }

        // XOR entropy with existing pad data
        for (size_t i = 0; i < (size_t)entropy_read; i++) {
            pad_data[i] ^= buffer[i];
        }

        // Seek back and write modified data
        if (fseek(pad_file, bytes_processed, SEEK_SET) != 0) {
            if (display_progress) {
                printf("Error: Cannot seek back to position %zu in pad file\n", bytes_processed);
            }
            error_occurred = 1;
            break;
        }

        if (fwrite(pad_data, 1, entropy_read, pad_file) != (size_t)entropy_read) {
            if (display_progress) {
                printf("Error: Cannot write modified pad data\n");
            }
            error_occurred = 1;
            break;
        }

        bytes_processed += entropy_read;

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

    close(device_fd);
    fclose(pad_file);

    if (error_occurred) {
        return 1;
    }

    if (display_progress) {
        show_progress(total_bytes, total_bytes, start_time);
        printf("\n✓ Successfully streamed %zu bytes of hardware entropy to pad\n", bytes_processed);
    }

    return 0;
}

// Detect all available hardware RNG devices
int detect_all_hardware_rng_devices(hardware_rng_device_t* devices, int max_devices, int* num_devices_found) {
    *num_devices_found = 0;

    // Scan /dev directory for serial devices (ttyUSB*, ttyACM*)
    DIR* dev_dir = opendir("/dev");
    if (!dev_dir) {
        return 1; // Error opening /dev
    }

    struct dirent* entry;
    while ((entry = readdir(dev_dir)) != NULL && *num_devices_found < max_devices) {
        // Check for serial device patterns
        if (strncmp(entry->d_name, "ttyUSB", 6) == 0 || strncmp(entry->d_name, "ttyACM", 6) == 0) {
            char device_path[512]; // Increased buffer size to prevent truncation
            int ret = snprintf(device_path, sizeof(device_path), "/dev/%s", entry->d_name);
            if (ret >= (int)sizeof(device_path)) {
                continue; // Skip if path would be truncated
            }

            // Check if this is a TrueRNG/SwiftRNG device by reading VID/PID
            char vid[5], pid[5];
            if (read_usb_device_info(device_path, vid, pid) == 0) {
                hardware_rng_device_type_t device_type = 0;

                // Check against known TrueRNG VID/PID combinations
                if (strcmp(vid, TRUERNG_VID) == 0 && strcmp(pid, TRUERNG_ORIGINAL_PID) == 0) {
                    device_type = TRUERNG_ORIGINAL;
                } else if (strcmp(vid, TRUERNG_VID) == 0 && strcmp(pid, TRUERNG_PRO_PID) == 0) {
                    device_type = TRUERNG_PRO;
                } else if (strcmp(vid, TRUERNG_VID) == 0 && strcmp(pid, TRUERNG_PRO_V2_PID) == 0) {
                    device_type = TRUERNG_PRO_V2;
                }

                if (device_type != 0) {
                    // Found a TrueRNG/SwiftRNG device
                    hardware_rng_device_t* device = &devices[*num_devices_found];

                    strncpy(device->port_path, device_path, sizeof(device->port_path) - 1);
                    device->device_type = device_type;
                    strncpy(device->friendly_name, get_truerng_device_name(device_type), sizeof(device->friendly_name) - 1);

                    // Assume device is working if VID/PID matches (no test needed)
                    device->is_working = 1;

                    (*num_devices_found)++;
                }
            }
        }
    }

    closedir(dev_dir);
    return 0; // Success
}

// Configure serial port for different RNG device types
int configure_rng_serial_port(int fd, hardware_rng_device_type_t device_type) {
    (void)device_type; // Suppress unused parameter warning - all TrueRNG devices use same config
    struct termios tty;

    if (tcgetattr(fd, &tty) != 0) {
        return 1; // Error getting terminal attributes
    }

    // TrueRNG configuration - traditional serial settings
    // TrueRNG devices: 115200 baud, 8N1, no flow control
    cfsetospeed(&tty, B115200);
    cfsetispeed(&tty, B115200);
    tty.c_cflag = (tty.c_cflag & ~CSIZE) | CS8;     // 8-bit chars
    tty.c_cflag |= CLOCAL | CREAD;                   // ignore modem controls, enable reading
    tty.c_cflag &= ~(PARENB | PARODD);              // no parity
    tty.c_cflag &= ~CSTOPB;                          // 1 stop bit
    tty.c_cflag &= ~CRTSCTS;                         // no hardware flow control
    tty.c_iflag &= ~(IXON | IXOFF | IXANY);          // no software flow control
    tty.c_iflag &= ~(ICANON | ECHO | ECHOE | ISIG);  // raw mode
    tty.c_oflag &= ~OPOST;                           // raw output

    // Set timeouts for TrueRNG
    tty.c_cc[VMIN] = 1;   // read at least 1 character
    tty.c_cc[VTIME] = 10; // 1 second timeout

    if (tcsetattr(fd, TCSANOW, &tty) != 0) {
        return 1; // Error setting terminal attributes
    }

    // Flush any existing data
    tcflush(fd, TCIOFLUSH);

    return 0; // Success
}

// Collect entropy from a specific TrueRNG device
int collect_truerng_entropy_from_device(const hardware_rng_device_t* device, unsigned char* entropy_buffer,
                                       size_t target_bytes, size_t* collected_bytes, int display_progress) {
    if (!device || !entropy_buffer || !collected_bytes || target_bytes == 0) {
        return 1; // Invalid parameters
    }


    // Open the TrueRNG device
    int device_fd = open(device->port_path, O_RDONLY | O_NOCTTY);
    if (device_fd < 0) {
        if (display_progress) {
            printf("Error: Cannot open RNG device %s: %s\n", device->port_path, strerror(errno));
        }
        return 1;
    }

    // Configure serial port for this device type
    if (configure_rng_serial_port(device_fd, device->device_type) != 0) {
        if (display_progress) {
            printf("Error: Failed to configure serial port for %s\n", device->friendly_name);
        }
        close(device_fd);
        return 1;
    }

    // Standard delay for TrueRNG devices
    usleep(100000); // 100ms

    if (display_progress) {
        printf("Collecting %zu bytes from %s...\n", target_bytes, device->friendly_name);
    }

    // Read entropy data with timeout protection
    size_t total_read = 0;
    time_t start_time = time(NULL);
    time_t last_progress_time = start_time;

    while (total_read < target_bytes) {
        // Check for overall timeout (5 minutes max for large collections)
        time_t current_time = time(NULL);
        if (difftime(current_time, start_time) > 300) { // 5 minutes timeout
            if (display_progress) {
                printf("Error: Collection timeout - device may be unresponsive\n");
            }
            close(device_fd);
            return 1;
        }

        size_t remaining = target_bytes - total_read;
        size_t chunk_size = (remaining > 4096) ? 4096 : remaining; // Read in 4KB chunks

        ssize_t bytes_read = read(device_fd, entropy_buffer + total_read, chunk_size);
        if (bytes_read < 0) {
            if (errno == EINTR) {
                continue; // Interrupted, try again
            }
            if (errno == EAGAIN || errno == EWOULDBLOCK) {
                // Timeout occurred, check if we have enough data for a test
                if (total_read > 0 && target_bytes > 1024) {
                    // For testing purposes, we have enough data
                    break;
                }
                // For small collections, this is an error
                if (display_progress) {
                    printf("Error: Device read timeout - no data received\n");
                }
                close(device_fd);
                return 1;
            }
            if (display_progress) {
                printf("Error: Failed to read from TrueRNG device: %s\n", strerror(errno));
                printf("Device may have been disconnected.\n");
            }
            close(device_fd);
            return 1;
        }

        if (bytes_read == 0) {
            // End of data - this shouldn't happen for RNG devices
            if (total_read == 0) {
                if (display_progress) {
                    printf("Error: TrueRNG device returned no data (device disconnected or misconfigured?)\n");
                }
                close(device_fd);
                return 1;
            } else {
                // We have some data, might be enough for testing
                break;
            }
        }

        total_read += bytes_read;

        // Show progress
        if (display_progress && (total_read % 1024 == 0 || difftime(current_time, last_progress_time) >= 1)) {
            show_progress(total_read, target_bytes, start_time);
            last_progress_time = current_time;
        }
    }

    close(device_fd);

    if (display_progress) {
        show_progress(target_bytes, target_bytes, start_time);
        printf("\n✓ Successfully collected %zu bytes from TrueRNG device\n", total_read);
    }

    *collected_bytes = total_read;
    return 0;
}

// Read USB device VID/PID information from sysfs
int read_usb_device_info(const char* device_path, char* vid, char* pid) {
    // Extract device name from path (e.g., /dev/ttyUSB0 -> ttyUSB0)
    const char* device_name = strrchr(device_path, '/');
    if (!device_name) device_name = device_path;
    else device_name++; // Skip the '/'

    // Construct sysfs path for USB device info
    char sysfs_path[256];
    snprintf(sysfs_path, sizeof(sysfs_path), "/sys/class/tty/%s/device/../idVendor", device_name);

    FILE* vid_file = fopen(sysfs_path, "r");
    if (!vid_file) {
        return 1; // Cannot read VID
    }

    if (fscanf(vid_file, "%4s", vid) != 1) {
        fclose(vid_file);
        return 1; // Cannot parse VID
    }
    fclose(vid_file);

    // Read PID
    snprintf(sysfs_path, sizeof(sysfs_path), "/sys/class/tty/%s/device/../idProduct", device_name);
    FILE* pid_file = fopen(sysfs_path, "r");
    if (!pid_file) {
        return 1; // Cannot read PID
    }

    if (fscanf(pid_file, "%4s", pid) != 1) {
        fclose(pid_file);
        return 1; // Cannot parse PID
    }
    fclose(pid_file);

    return 0; // Success
}

// Get friendly name for hardware RNG device type
const char* get_truerng_device_name(hardware_rng_device_type_t device_type) {
    switch (device_type) {
        case TRUERNG_ORIGINAL:
            return "TrueRNG";
        case TRUERNG_PRO:
            return "TrueRNG Pro";
        case TRUERNG_PRO_V2:
            return "TrueRNG Pro V2";
        default:
            return "Unknown Hardware RNG Device";
    }
}

// Test if a hardware RNG device is working by attempting to read from it
int test_hardware_rng_device(const hardware_rng_device_t* device) {
    int fd = open(device->port_path, O_RDONLY | O_NONBLOCK);
    if (fd < 0) {
        return 1; // Cannot open device
    }

    // Try to read a small amount of data
    unsigned char test_buffer[16];
    ssize_t bytes_read = read(fd, test_buffer, sizeof(test_buffer));

    close(fd);

    if (bytes_read <= 0) {
        return 1; // Cannot read from device
    }

    return 0; // Device appears to be working
}

// Interactive device selection for hardware RNG
int select_hardware_rng_device_interactive(hardware_rng_device_t* devices, int num_devices, hardware_rng_device_t* selected_device) {
    if (num_devices == 0) {
        printf("No hardware RNG devices found.\n");
        return 1; // No devices available
    }

    if (num_devices == 1) {
        // Only one device, use it automatically
        *selected_device = devices[0];
        printf("Using %s (%s)\n\n", devices[0].friendly_name, devices[0].port_path);
        return 0;
    }

    // Multiple devices - let user choose
    printf("\nAvailable Hardware RNG Devices:\n");
    for (int i = 0; i < num_devices; i++) {
        printf("%d. %s (%s)\n",
               i + 1,
               devices[i].friendly_name,
               devices[i].port_path);
    }

    printf("\nSelect device (1-%d): ", num_devices);

    char input[10];
    if (fgets(input, sizeof(input), stdin) == NULL) {
        return 1; // Input error
    }

    int choice = atoi(input);
    if (choice < 1 || choice > num_devices) {
        printf("Invalid selection.\n");
        return 1;
    }

    *selected_device = devices[choice - 1];
    printf("Selected: %s (%s)\n", selected_device->friendly_name, selected_device->port_path);
    return 0;
}