/*
 * TrueRNG Read - C Implementation
 * Equivalent to main.py
 *
 * Requires: gcc, Linux system with termios support
 * Compile with: gcc -o truerng main.c
 *
 * On Linux - may need to be root or set /dev/tty port permissions to 666
 */

#define _GNU_SOURCE
#define _POSIX_C_SOURCE 200809L

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <fcntl.h>
#include <termios.h>
#include <time.h>
#include <sys/time.h>
#include <dirent.h>
#include <errno.h>
#include <ctype.h>
#include <sys/ioctl.h>
#include <getopt.h>
#include <limits.h>

// Define CRTSCTS if not available
#ifndef CRTSCTS
#define CRTSCTS 020000000000
#endif

// Constants
#define BLOCKSIZE 1024
#define NUMLOOPS 1024
#define MAX_PATH 512
#define MAX_PORTS 20

// TrueRNG Device IDs
#define TRUERNG_PID "04D8"
#define TRUERNG_HID "F5FE"
#define TRUERNGPRO_PID "16D0"
#define TRUERNGPRO_HID "0AA0"
#define TRUERNGPROV2_PID "04D8"
#define TRUERNGPROV2_HID "EBB5"

// Output formats
enum Format { BINARY, HEX, BASE64, DECIMAL };

// Command-line options structure
struct Options {
    long long bytes;
    enum Format format;
    char *output_file;
    int quiet;
    int verbose;
    int help;
};

// Global options
struct Options options;

// Global variables
static int ones_in_byte[256];

// Parse byte size with optional suffix (K, MB, GB, TB)
// Returns: parsed size in bytes, or -1 on error
// Sets *error to 1 if parsing fails, 0 on success
long long parse_byte_size(const char* size_str, int* error) {
    if (!size_str || !error) {
        if (error) *error = 1;
        return -1;
    }
    
    *error = 0;
    char* endptr;
    double value = strtod(size_str, &endptr);
    
    if (value < 0) {
        *error = 1;
        return -1;
    }
    
    // If no suffix, return the value as bytes (backward compatibility)
    if (*endptr == '\0') {
        if (value > LLONG_MAX) {
            *error = 1;
            return -1;
        }
        return (long long)value;
    }
    
    // Parse suffix (case-insensitive)
    char suffix[4] = {0};
    int i = 0;
    while (*endptr && i < 3) {
        suffix[i++] = toupper(*endptr++);
    }
    
    // Determine multiplier
    long long multiplier = 1;
    if (strcmp(suffix, "K") == 0 || strcmp(suffix, "KB") == 0) {
        multiplier = 1024LL;
    } else if (strcmp(suffix, "M") == 0 || strcmp(suffix, "MB") == 0) {
        multiplier = 1024LL * 1024LL;
    } else if (strcmp(suffix, "G") == 0 || strcmp(suffix, "GB") == 0) {
        multiplier = 1024LL * 1024LL * 1024LL;
    } else if (strcmp(suffix, "T") == 0 || strcmp(suffix, "TB") == 0) {
        multiplier = 1024LL * 1024LL * 1024LL * 1024LL;
    } else {
        *error = 1;
        return -1;
    }
    
    // Check for overflow
    if (value > (double)LLONG_MAX / multiplier) {
        *error = 1;
        return -1;
    }
    
    return (long long)(value * multiplier);
}

// Output format functions
void output_binary(unsigned char *buffer, int bytes_read, FILE *out) {
    fwrite(buffer, 1, bytes_read, out);
}

void output_hex(unsigned char *buffer, int bytes_read, FILE *out, int piped) {
    (void)piped; // Suppress unused parameter warning
    for (int i = 0; i < bytes_read; i++) {
        fprintf(out, "%02x", buffer[i]);
    }
    // Always add newline for text formats (standard Linux practice)
    fprintf(out, "\n");
}

void output_decimal(unsigned char *buffer, int bytes_read, FILE *out, int piped) {
    (void)piped; // Suppress unused parameter warning
    for (int i = 0; i < bytes_read; i++) {
        fprintf(out, "%d", buffer[i]);
    }
    // Always add newline for text formats (standard Linux practice)
    fprintf(out, "\n");
}

void output_base64(unsigned char *buffer, int bytes_read, FILE *out, int piped) {
    (void)piped; // Suppress unused parameter warning
    static const char b64[] = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/";
    char encoded[5];
    for (int i = 0; i < bytes_read; i += 3) {
        encoded[0] = b64[buffer[i] >> 2];
        
        if (i + 1 < bytes_read) {
            encoded[1] = b64[((buffer[i] & 0x03) << 4) | (buffer[i+1] >> 4)];
        } else {
            encoded[1] = b64[(buffer[i] & 0x03) << 4];
        }
        
        if (i + 2 < bytes_read) {
            encoded[2] = b64[((buffer[i+1] & 0x0f) << 2) | (buffer[i+2] >> 6)];
            encoded[3] = b64[buffer[i+2] & 0x3f];
        } else if (i + 1 < bytes_read) {
            encoded[2] = b64[(buffer[i+1] & 0x0f) << 2];
            encoded[3] = '=';
        } else {
            encoded[2] = '=';
            encoded[3] = '=';
        }
        
        fwrite(encoded, 1, 4, out);
    }
    // Always add newline for text formats (standard Linux practice)
    fprintf(out, "\n");
}

void print_usage(void) {
    printf("TrueRNG - True Random Number Generator\n");
    printf("Usage: truerng [OPTIONS]\n\n");
    printf("Options:\n");
    printf("  -n, --bytes <N>          Number of bytes to generate (default: 1048576)\n");
    printf("                           Supports suffixes: K, MB, GB, TB (e.g., 1K, 2.5MB, 1GB)\n");
    printf("  -f, --format <FORMAT>    Output format: binary, hex, base64, decimal (default: binary when piped, interactive otherwise)\n");
    printf("  -o, --output <FILE>      Output filename (ignored in piped mode)\n");
    printf("  -q, --quiet              Suppress statistics/progress\n");
    printf("  -v, --verbose            Show detailed device information\n");
    printf("  -h, --help               Show this help message\n");
    printf("\nExamples:\n");
    printf("  truerng -n 1024 -f hex                    # Interactive mode with hex output\n");
    printf("  truerng -n 1K -f hex                      # Same as above using K suffix\n");
    printf("  truerng -n 2.5MB | xxd                    # Piped mode with MB suffix\n");
    printf("  truerng -n 1GB -o random.dat -q           # Save 1GB to file quietly\n");
    printf("  truerng -n 512K -f hex -o output.hex      # Save 512KB as hex to file\n");
}

int parse_arguments(int argc, char *argv[]) {
    static const char *opt_string = "n:f:o:qvh";
    static struct option long_options[] = {
        {"bytes", required_argument, 0, 'n'},
        {"format", required_argument, 0, 'f'},
        {"output", required_argument, 0, 'o'},
        {"quiet", no_argument, 0, 'q'},
        {"verbose", no_argument, 0, 'v'},
        {"help", no_argument, 0, 'h'},
        {0, 0, 0, 0}
    };

    options.bytes = 1048576; // Default 1MB
    options.format = BINARY;
    options.output_file = NULL;
    options.quiet = 0;
    options.verbose = 0;
    options.help = 0;

    int opt;
    int option_index = 0;
    while ((opt = getopt_long(argc, argv, opt_string, long_options, &option_index)) != -1) {
        switch (opt) {
            case 'n': {
                int parse_error;
                options.bytes = parse_byte_size(optarg, &parse_error);
                if (parse_error || options.bytes <= 0) {
                    fprintf(stderr, "Error: Invalid byte size '%s'. Use a positive number optionally followed by K, MB, GB, or TB\n", optarg);
                    return -1;
                }
                break;
            }
            case 'f':
                if (strcmp(optarg, "binary") == 0) options.format = BINARY;
                else if (strcmp(optarg, "hex") == 0) options.format = HEX;
                else if (strcmp(optarg, "base64") == 0) options.format = BASE64;
                else if (strcmp(optarg, "decimal") == 0) options.format = DECIMAL;
                else {
                    fprintf(stderr, "Error: Invalid format '%s'\n", optarg);
                    return -1;
                }
                break;
            case 'o':
                options.output_file = optarg;
                break;
            case 'q':
                options.quiet = 1;
                break;
            case 'v':
                options.verbose = 1;
                break;
            case 'h':
                options.help = 1;
                return 0;
            default:
                return -1;
        }
    }

    return 0;
}

int is_piped(void) {
    return !isatty(STDOUT_FILENO);
}

// Function prototypes
void init_ones_lookup_table(void);
int find_truerng_port(char* port_path, int piped, int verbose);
int setup_serial_port(const char* port_path);
int read_usb_device_info(const char* port_name, char* vid, char* pid);
double get_time_diff(struct timeval start, struct timeval end);
long long parse_byte_size(const char* size_str, int* error);
int parse_arguments(int argc, char *argv[]);
void print_usage(void);
int is_piped(void);
void output_binary(unsigned char *buffer, int bytes_read, FILE *out);
void output_hex(unsigned char *buffer, int bytes_read, FILE *out, int piped);
void output_decimal(unsigned char *buffer, int bytes_read, FILE *out, int piped);
void output_base64(unsigned char *buffer, int bytes_read, FILE *out, int piped);

// Initialize lookup table for counting ones in a byte
void init_ones_lookup_table(void) {
    int i, count;
    for (i = 0; i < 256; i++) {
        count = 0;
        int temp = i;
        while (temp) {
            count += temp & 1;
            temp >>= 1;
        }
        ones_in_byte[i] = count;
    }
}

// Read USB device info from sysfs
int read_usb_device_info(const char* port_name, char* vid, char* pid) {
    char path[MAX_PATH];
    FILE *fp;
    
    // Try to read idVendor first (works for both ttyUSB and ttyACM devices)
    snprintf(path, sizeof(path), "/sys/class/tty/%s/device/../idVendor", port_name);
    fp = fopen(path, "r");
    if (fp) {
        if (fgets(vid, 8, fp) != NULL) {
            // Remove newline if present
            int len = strlen(vid);
            if (len > 0 && vid[len-1] == '\n') {
                vid[len-1] = '\0';
            }
        } else {
            fclose(fp);
            return 0;
        }
        fclose(fp);
    } else {
        return 0;
    }
    
    // Try to read idProduct
    snprintf(path, sizeof(path), "/sys/class/tty/%s/device/../idProduct", port_name);
    fp = fopen(path, "r");
    if (fp) {
        if (fgets(pid, 8, fp) != NULL) {
            // Remove newline if present
            int len = strlen(pid);
            if (len > 0 && pid[len-1] == '\n') {
                pid[len-1] = '\0';
            }
        } else {
            fclose(fp);
            return 0;
        }
        fclose(fp);
        return 1;
    } else {
        return 0;
    }
}

// Find TrueRNG device port
// Returns: 0=failure, 1=TrueRNGproV2, 2=TrueRNGpro, 3=TrueRNG
int find_truerng_port(char* port_path, int piped, int verbose) {
    DIR *dir;
    struct dirent *entry;
    char vid[8], pid[8];
    int device_found = 0;
    
    dir = opendir("/dev");
    if (dir == NULL) {
        perror("Cannot open /dev directory");
        return 0;
    }
    
    while ((entry = readdir(dir)) != NULL) {
        // Look for ttyUSB* or ttyACM* devices
        if (strncmp(entry->d_name, "ttyUSB", 6) == 0 || 
            strncmp(entry->d_name, "ttyACM", 6) == 0) {
            
            if (read_usb_device_info(entry->d_name, vid, pid)) {
                // Convert to uppercase for comparison
                for (int i = 0; vid[i]; i++) vid[i] = toupper(vid[i]);
                for (int i = 0; pid[i]; i++) pid[i] = toupper(pid[i]);
                
                // Check for TrueRNGproV2
                if (strcmp(vid, TRUERNGPROV2_PID) == 0 && strcmp(pid, TRUERNGPROV2_HID) == 0) {
                    snprintf(port_path, MAX_PATH, "/dev/%s", entry->d_name);
                    if (!piped || verbose) {
                        printf("TrueRNGproV2 Found\n");
                    }
                    device_found = 1;
                    break;
                }

                // Check for TrueRNGpro
                if (strcmp(vid, TRUERNGPRO_PID) == 0 && strcmp(pid, TRUERNGPRO_HID) == 0) {
                    snprintf(port_path, MAX_PATH, "/dev/%s", entry->d_name);
                    if (!piped || verbose) {
                        printf("TrueRNGpro Found\n");
                    }
                    device_found = 2;
                    break;
                }

                // Check for TrueRNG
                if (strcmp(vid, TRUERNG_PID) == 0 && strcmp(pid, TRUERNG_HID) == 0) {
                    snprintf(port_path, MAX_PATH, "/dev/%s", entry->d_name);
                    if (!piped || verbose) {
                        printf("TrueRNG Found\n");
                    }
                    device_found = 3;
                    break;
                }
            }
        }
    }
    
    closedir(dir);
    return device_found;
}

// Setup serial port
int setup_serial_port(const char* port_path) {
    int fd;
    struct termios tty;
    
    fd = open(port_path, O_RDWR | O_NOCTTY);
    if (fd < 0) {
        printf("Port Not Usable!\n");
        printf("Do you have permissions set to read %s ?\n", port_path);
        return -1;
    }
    
    // Get current port settings
    if (tcgetattr(fd, &tty) != 0) {
        printf("Error getting port attributes\n");
        close(fd);
        return -1;
    }
    
    // Set baud rate (most TrueRNG devices use default settings)
    cfsetospeed(&tty, B9600);
    cfsetispeed(&tty, B9600);
    
    // 8N1 mode
    tty.c_cflag &= ~PARENB;   // No parity
    tty.c_cflag &= ~CSTOPB;   // 1 stop bit
    tty.c_cflag &= ~CSIZE;    // Clear size bits
    tty.c_cflag |= CS8;       // 8 data bits
    tty.c_cflag &= ~CRTSCTS;  // No hardware flow control
    tty.c_cflag |= CREAD | CLOCAL; // Enable reading and ignore modem controls
    
    // Raw input mode
    tty.c_lflag &= ~(ICANON | ECHO | ECHOE | ISIG);
    tty.c_iflag &= ~(IXON | IXOFF | IXANY);
    tty.c_oflag &= ~OPOST;
    
    // Set for blocking reads - wait for data indefinitely
    tty.c_cc[VMIN] = 1;   // Block until at least 1 character is received
    tty.c_cc[VTIME] = 0;  // No timeout
    
    // Apply settings
    if (tcsetattr(fd, TCSANOW, &tty) != 0) {
        printf("Error setting port attributes\n");
        close(fd);
        return -1;
    }
    
    // Flush input buffer
    tcflush(fd, TCIFLUSH);
    
    // Set DTR
    int status;
    ioctl(fd, TIOCMGET, &status);
    status |= TIOCM_DTR;
    ioctl(fd, TIOCMSET, &status);
    
    return fd;
}

// Calculate time difference in seconds
double get_time_diff(struct timeval start, struct timeval end) {
    return (end.tv_sec - start.tv_sec) + (end.tv_usec - start.tv_usec) / 1000000.0;
}

int main(int argc, char *argv[]) {
    // Parse arguments first
    if (parse_arguments(argc, argv)) {
        return 1;
    }

    // Handle help
    if (options.help) {
        print_usage();
        return 0;
    }

    // Determine mode: interactive (no args) vs piped (has args or actual piping)
    int interactive_mode = (argc == 1);
    int piped_mode = !interactive_mode || is_piped();

    // Set default format based on mode
    if (interactive_mode && options.format == BINARY) {
        // Interactive mode defaults to hex for readability
        options.format = HEX;
    }

    // Print headers based on mode
    if (interactive_mode && !options.quiet) {
        // Interactive mode: always show headers unless quiet
        printf("TrueRNG - True Random Number Generator\n");
        printf("==================================================\n");
    } else if (piped_mode && options.verbose && !options.quiet) {
        // Piped mode: only show headers if verbose
        printf("TrueRNG - True Random Number Generator\n");
        printf("Mode: Piped (Verbose)\n");
        printf("==================================================\n");
    }

    char port_path[MAX_PATH];
    int serial_fd;
    FILE *fp;
    unsigned char buffer[BLOCKSIZE];
    long long totalbytes = 0;
    long long totalones = 0;
    long long totalzeros = 0;
    struct timeval starttime, currenttime;
    double elapsed_time, rate;

    // Initialize lookup table
    init_ones_lookup_table();

    // Find TrueRNG device
    if (!find_truerng_port(port_path, piped_mode, options.verbose)) {
        fprintf(stderr, "TrueRNG Not Found\n");
        return 1;
    }

    // Setup output
    if (piped_mode && !options.output_file) {
        fp = stdout;
    } else if (options.output_file) {
        fp = fopen(options.output_file, "w");
        if (fp == NULL) {
            fprintf(stderr, "Error Opening File: %s\n", options.output_file);
            return 1;
        }
    } else {
        fp = stdout;
    }

    // Setup serial port
    serial_fd = setup_serial_port(port_path);
    if (serial_fd < 0) {
        if (options.output_file && fp != stdout) fclose(fp);
        return 1;
    }

    if (interactive_mode && !options.quiet) {
        printf("Using port: %s\n", port_path);
        printf("Generating: %lld bytes\n", options.bytes);
        if (options.output_file) {
            printf("Output file: %s\n", options.output_file);
        }
        printf("Starting data collection...\n");
    } else if (piped_mode && options.verbose && !options.quiet) {
        printf("Using port: %s\n", port_path);
        printf("Generating: %lld bytes\n", options.bytes);
        printf("Starting data collection...\n");
    }

    // Get start time
    gettimeofday(&starttime, NULL);

    // Main loop - read until we have enough bytes
    while (totalbytes < options.bytes) {
        long long remaining = options.bytes - totalbytes;
        int blocksize = (remaining < BLOCKSIZE) ? remaining : BLOCKSIZE;

        int bytes_in_buffer = 0;
        while (bytes_in_buffer < blocksize) {
            ssize_t bytes_read = read(serial_fd, buffer + bytes_in_buffer, blocksize - bytes_in_buffer);
            if (bytes_read < 0) {
                fprintf(stderr, "Read Failed!!! Error: %s\n", strerror(errno));
                goto cleanup;
            } else if (bytes_read == 0) {
                fprintf(stderr, "\nNo data received, device may have stopped\n");
                goto cleanup;
            }
            bytes_in_buffer += bytes_read;
        }

        // Update total bytes
        totalbytes += blocksize;

        // Count ones and zeros
        for (int i = 0; i < blocksize; i++) {
            int ones = ones_in_byte[buffer[i]];
            totalones += ones;
            totalzeros += (8 - ones);
        }

        // Output data
        switch (options.format) {
            case BINARY:
                output_binary(buffer, blocksize, fp);
                break;
            case HEX:
                output_hex(buffer, blocksize, fp, piped_mode);
                break;
            case BASE64:
                output_base64(buffer, blocksize, fp, piped_mode);
                break;
            case DECIMAL:
                output_decimal(buffer, blocksize, fp, piped_mode);
                break;
        }

        // Display progress only in interactive mode (and verbose piped mode)
        gettimeofday(&currenttime, NULL);
        elapsed_time = get_time_diff(starttime, currenttime);
        if (interactive_mode && !options.quiet) {
            if (elapsed_time > 0) {
                rate = (double)totalbytes / (elapsed_time * 1000.0);
                printf("%lld/%lld Bytes Read at %6.2f Kbytes/s\r",
                        totalbytes, options.bytes, rate);
                fflush(stdout);
            }
        } else if (piped_mode && options.verbose && !options.quiet && elapsed_time > 0) {
            rate = (double)totalbytes / (elapsed_time * 1000.0);
            printf("\r%lld/%lld Bytes Read at %6.2f Kbytes/s",
                    totalbytes, options.bytes, rate);
            fflush(stdout);
        }
    }

cleanup:

    // Print results only in interactive mode or piped verbose mode
    if ((interactive_mode && !options.quiet) || (piped_mode && options.verbose)) {
        gettimeofday(&currenttime, NULL);
        double total_time = get_time_diff(starttime, currenttime);

        printf("\n\nResults\n");
        printf("=======\n");
        printf("Total time: %.2f seconds\n", total_time);
        printf("Total Ones:  %lld\n", totalones);
        printf("Total Zeros: %lld\n", totalzeros);
        printf("Total Bits:  %lld\n", totalbytes * 8);

        if (totalones == totalzeros) {
            printf("Equal number of ones and zeros\n");
        } else if (totalones > totalzeros) {
            printf("Extra ones:  %lld\n", totalones - totalzeros);
        } else {
            printf("Extra zeros: %lld\n", totalzeros - totalones);
        }
        printf("=======\n");
    }

    // Cleanup
    close(serial_fd);
    if (options.output_file && fp != stdout) fclose(fp);

    // Reset terminal min setting
    char stty_cmd[MAX_PATH + 20];
    snprintf(stty_cmd, sizeof(stty_cmd), "stty -F %s min 1", port_path);
    if (system(stty_cmd) != 0) {
        // Ignore failure - this is just cleanup
    }

    return 0;
}