#include <stdio.h>
#include <stdlib.h>
#include <unistd.h>
#include <dlfcn.h>
#include <sys/types.h>
#include <sys/stat.h>
#include <fcntl.h>
#include <time.h>
#include <sys/uio.h>
#include <sys/sendfile.h>
#include <linux/fs.h>
#include <string.h>
#include <stdarg.h>
#include <sys/socket.h>

// Function pointers for original read and write and other relevant calls

// Function pointers for standard I/O functions
int (*original_putchar_unlocked)(int);
int (*original_putc)(int, FILE *);
int (*original_putc_unlocked)(int, FILE *);
size_t (*original_fread)(void *, size_t, size_t, FILE *);
size_t (*original_fread_unlocked)(void *, size_t, size_t, FILE *);
int (*original_fgetc)(FILE *);
char *(*original_fgets)(char *, int, FILE *);
int (*original_getc)(FILE *);
int (*original_getc_unlocked)(FILE *);

// Function pointers for socket functions
ssize_t (*original_recv)(int, void *, size_t, int);
ssize_t (*original_recvfrom)(int, void *, size_t, int, struct sockaddr *, socklen_t *);
ssize_t (*original_recvmsg)(int, struct msghdr *, int);

// Function pointers for formatted input functions
int (*original_scanf)(const char *, ...);
int (*original_fscanf)(FILE *, const char *, ...);

// Additional low-level function pointers
ssize_t (*original_pread)(int fd, void *buf, size_t count, off_t offset);
ssize_t (*original_pwrite)(int fd, const void *buf, size_t count, off_t offset);
ssize_t (*original_readv)(int fd, const struct iovec *iov, int iovcnt);
ssize_t (*original_writev)(int fd, const struct iovec *iov, int iovcnt);


ssize_t (*original_read)(int fd, void *buf, size_t count);
ssize_t (*original_write)(int fd, const void *buf, size_t count);
size_t (*original_fwrite)(const void *ptr, size_t size, size_t nmemb, FILE *stream);
size_t (*original_fwrite_unlocked)(const void *ptr, size_t size, size_t nmemb, FILE *stream);
int (*original_fflush)(FILE *stream);
int (*original_fputs)(const char *str, FILE *stream);
int (*original_fputs_unlocked)(const char *str, FILE *stream);
int (*original_putchar)(int c);
int (*original_dprintf)(int fd, const char *format, ...);
ssize_t (*original_pread)(int fd, void *buf, size_t count, off_t offset);
ssize_t (*original_pwrite)(int fd, const void *buf, size_t count, off_t offset);
ssize_t (*original_readv)(int fd, const struct iovec *iov, int iovcnt);
ssize_t (*original_writev)(int fd, const struct iovec *iov, int iovcnt);
off_t (*original_sendfile)(int out_fd, int in_fd, off_t *offset, size_t count);
ssize_t (*original_splice)(int fd_in, loff_t *off_in, int fd_out, loff_t *off_out, size_t len, unsigned int flags);
ssize_t (*original_copy_file_range)(int fd_in, loff_t *off_in, int fd_out, loff_t *off_out, size_t len, unsigned int flags);

// Function pointer for original functions
ssize_t (*original_send)(int sockfd, const void *buf, size_t len, int flags);
ssize_t (*original_sendto)(int sockfd, const void *buf, size_t len, int flags,
                            const struct sockaddr *dest_addr, socklen_t addrlen);
ssize_t (*original_sendmsg)(int sockfd, const struct msghdr *msg, int flags);
ssize_t (*original_vmsplice)(int fd, const struct iovec *iov, unsigned long nr_segs,
                             unsigned int flags);
ssize_t (*original_process_vm_readv)(pid_t pid, const struct iovec *local_iov,
                                      int local_iovcnt, const struct iovec *remote_iov,
                                      int remote_iovcnt, unsigned long flags);
ssize_t (*original_process_vm_writev)(pid_t pid, const struct iovec *local_iov,
                                       int local_iovcnt, const struct iovec *remote_iov,
                                       int remote_iovcnt, unsigned long flags);

// Rate limits in bytes per second
static size_t read_limit = 0; // 0 = no limit
static size_t write_limit = 0;
static int debug_mode = 0;

// Tracking variables
static size_t read_bytes = 0;
static size_t write_bytes = 0;
static struct timespec read_start, write_start;

void throttle(size_t *bytes, size_t limit, struct timespec *start) {
    if (limit == 0) {
        return; // No limit applied
    }

    struct timespec now;
    clock_gettime(CLOCK_MONOTONIC, &now);

    double elapsed = (now.tv_sec - start->tv_sec) + (now.tv_nsec - start->tv_nsec) / 1e9;
    // If it is > 1 second since last write:
    // treat as if it is only 1 second
    elapsed = elapsed > 1 ? 1 : elapsed;
    double allowed_bytes = limit * elapsed;
    if (debug_mode) fprintf(stderr, "allowed: %ld written: %ld\n",
                            (long)allowed_bytes, *bytes);
    if (*bytes > allowed_bytes) {
        *bytes -= allowed_bytes;
        double sleep_time = *bytes / (double)limit;

        if (debug_mode) fprintf(stderr, "sleep: %f\n", sleep_time);
        usleep((useconds_t)(sleep_time * 1e6));
        *start = now;
    }
}

size_t fwrite(const void *ptr, size_t size, size_t nmemb, FILE *stream) {
    if (!original_fwrite) {
        original_fwrite = dlsym(RTLD_NEXT, "fwrite");
        if (!original_fwrite) {
            if (debug_mode) fprintf(stderr, "Error loading original fwrite: %s\n", dlerror());
            exit(EXIT_FAILURE);
        }
    }

    size_t total_bytes = size * nmemb;
    throttle(&write_bytes, write_limit, &write_start);

    size_t result = original_fwrite(ptr, size, nmemb, stream);
    if (result > 0) {
        write_bytes += total_bytes;
    }
    return result;
}

#undef fwrite_unlocked
size_t fwrite_unlocked(const void *ptr, size_t size, size_t nmemb, FILE *stream) {
    if (!original_fwrite_unlocked) {
        original_fwrite_unlocked = dlsym(RTLD_NEXT, "fwrite_unlocked");
        if (!original_fwrite_unlocked) {
            if (debug_mode) fprintf(stderr, "Error loading original fwrite_unlocked: %s\n", dlerror());
            exit(EXIT_FAILURE);
        }
    }

    size_t total_bytes = size * nmemb;
    throttle(&write_bytes, write_limit, &write_start);

    size_t result = original_fwrite_unlocked(ptr, size, nmemb, stream);
    if (result > 0) {
        write_bytes += total_bytes;
    }
    return result;
}

int fputs(const char *str, FILE *stream) {
    if (!original_fputs) {
        original_fputs = dlsym(RTLD_NEXT, "fputs");
        if (!original_fputs) {
            if (debug_mode) fprintf(stderr, "Error loading original fputs: %s\n", dlerror());
            exit(EXIT_FAILURE);
        }
    }

    size_t total_bytes = strlen(str);
    throttle(&write_bytes, write_limit, &write_start);

    int result = original_fputs(str, stream);
    if (result >= 0) {
        write_bytes += total_bytes;
    }
    return result;
}

int fputs_unlocked(const char *str, FILE *stream) {
    if (!original_fputs_unlocked) {
        original_fputs_unlocked = dlsym(RTLD_NEXT, "fputs_unlocked");
        if (!original_fputs_unlocked) {
            if (debug_mode) fprintf(stderr, "Error loading original fputs_unlocked: %s\n", dlerror());
            exit(EXIT_FAILURE);
        }
    }

    size_t total_bytes = strlen(str);
    throttle(&write_bytes, write_limit, &write_start);

    int result = original_fputs_unlocked(str, stream);
    if (result >= 0) {
        write_bytes += total_bytes;
    }
    return result;
}

int putchar(int c) {
    if (!original_putchar) {
        original_putchar = dlsym(RTLD_NEXT, "putchar");
        if (!original_putchar) {
            if (debug_mode) fprintf(stderr, "Error loading original putchar: %s\n", dlerror());
            exit(EXIT_FAILURE);
        }
    }

    throttle(&write_bytes, write_limit, &write_start);

    int result = original_putchar(c);
    if (result != EOF) {
        write_bytes += 1;
    }
    return result;
}

int dprintf(int fd, const char *format, ...) {
    if (!original_dprintf) {
        original_dprintf = dlsym(RTLD_NEXT, "dprintf");
        if (!original_dprintf) {
            if (debug_mode) fprintf(stderr, "Error loading original dprintf: %s\n", dlerror());
            exit(EXIT_FAILURE);
        }
    }

    va_list args;
    va_start(args, format);
    char buffer[1024];
    int len = vsnprintf(buffer, sizeof(buffer), format, args);
    va_end(args);

    if (len > 0) {
        throttle(&write_bytes, write_limit, &write_start);
        write_bytes += len;
    }

    return original_dprintf(fd, "%s", buffer);
}

int fflush(FILE *stream) {
    if (!original_fflush) {
        original_fflush = dlsym(RTLD_NEXT, "fflush");
        if (!original_fflush) {
            if (debug_mode) fprintf(stderr, "Error loading original fflush: %s\n", dlerror());
            exit(EXIT_FAILURE);
        }
    }

    throttle(&write_bytes, write_limit, &write_start);

    return original_fflush(stream);
}

ssize_t read(int fd, void *buf, size_t count) {
    if (!original_read) {
        original_read = dlsym(RTLD_NEXT, "read");
        if (!original_read) {
            if (debug_mode) fprintf(stderr, "Error loading original read: %s\n", dlerror());
            exit(EXIT_FAILURE);
        }
    }

    throttle(&read_bytes, read_limit, &read_start);

    ssize_t result = original_read(fd, buf, count);
    if (result > 0) {
        read_bytes += result;
    }
    return result;
}

ssize_t write(int fd, const void *buf, size_t count) {
    if (!original_write) {
        original_write = dlsym(RTLD_NEXT, "write");
        if (!original_write) {
            if (debug_mode) fprintf(stderr, "Error loading original write: %s\n", dlerror());
            exit(EXIT_FAILURE);
        }
    }

    throttle(&write_bytes, write_limit, &write_start);

    ssize_t result = original_write(fd, buf, count);
    if (result > 0) {
        write_bytes += result;
    }
    return result;
}

ssize_t pread(int fd, void *buf, size_t count, off_t offset) {
    if (!original_pread) {
        original_pread = dlsym(RTLD_NEXT, "pread");
        if (!original_pread) {
            if (debug_mode) fprintf(stderr, "Error loading original pread: %s\n", dlerror());
            exit(EXIT_FAILURE);
        }
    }

    throttle(&read_bytes, read_limit, &read_start);

    ssize_t result = original_pread(fd, buf, count, offset);
    if (result > 0) {
        read_bytes += result;
    }
    return result;
}

ssize_t pwrite(int fd, const void *buf, size_t count, off_t offset) {
    if (!original_pwrite) {
        original_pwrite = dlsym(RTLD_NEXT, "pwrite");
        if (!original_pwrite) {
            if (debug_mode) fprintf(stderr, "Error loading original pwrite: %s\n", dlerror());
            exit(EXIT_FAILURE);
        }
    }

    throttle(&write_bytes, write_limit, &write_start);

    ssize_t result = original_pwrite(fd, buf, count, offset);
    if (result > 0) {
        write_bytes += result;
    }
    return result;
}

ssize_t readv(int fd, const struct iovec *iov, int iovcnt) {
    if (!original_readv) {
        original_readv = dlsym(RTLD_NEXT, "readv");
        if (!original_readv) {
            if (debug_mode) fprintf(stderr, "Error loading original readv: %s\n", dlerror());
            exit(EXIT_FAILURE);
        }
    }

    throttle(&read_bytes, read_limit, &read_start);

    ssize_t result = original_readv(fd, iov, iovcnt);
    if (result > 0) {
        read_bytes += result;
    }
    return result;
}

ssize_t writev(int fd, const struct iovec *iov, int iovcnt) {
    if (!original_writev) {
        original_writev = dlsym(RTLD_NEXT, "writev");
        if (!original_writev) {
            if (debug_mode) fprintf(stderr, "Error loading original writev: %s\n", dlerror());
            exit(EXIT_FAILURE);
        }
    }

    throttle(&write_bytes, write_limit, &write_start);

    ssize_t result = original_writev(fd, iov, iovcnt);
    if (result > 0) {
        write_bytes += result;
    }
    return result;
}

off_t sendfile(int out_fd, int in_fd, off_t *offset, size_t count) {
    if (!original_sendfile) {
        original_sendfile = dlsym(RTLD_NEXT, "sendfile");
        if (!original_sendfile) {
            if (debug_mode) fprintf(stderr, "Error loading original sendfile: %s\n", dlerror());
            exit(EXIT_FAILURE);
        }
    }

    throttle(&read_bytes, read_limit, &read_start);
    throttle(&write_bytes, write_limit, &write_start);

    off_t result = original_sendfile(out_fd, in_fd, offset, count);
    if (result > 0) {
        read_bytes += result;
        write_bytes += result;
    }
    return result;
}

ssize_t splice(int fd_in, loff_t *off_in, int fd_out, loff_t *off_out, size_t len, unsigned int flags) {
    if (!original_splice) {
        original_splice = dlsym(RTLD_NEXT, "splice");
        if (!original_splice) {
            if (debug_mode) fprintf(stderr, "Error loading original splice: %s\n", dlerror());
            exit(EXIT_FAILURE);
        }
    }

    throttle(&read_bytes, read_limit, &read_start);
    throttle(&write_bytes, write_limit, &write_start);

    ssize_t result = original_splice(fd_in, off_in, fd_out, off_out, len, flags);
    if (result > 0) {
        read_bytes += result;
        write_bytes += result;
    }
    return result;
}

ssize_t copy_file_range(int fd_in, loff_t *off_in, int fd_out, loff_t *off_out, size_t len, unsigned int flags) {
    if (!original_copy_file_range) {
        original_copy_file_range = dlsym(RTLD_NEXT, "copy_file_range");
        if (!original_copy_file_range) {
            if (debug_mode) fprintf(stderr, "Error loading original copy_file_range: %s\n", dlerror());
            exit(EXIT_FAILURE);
        }
    }

    throttle(&read_bytes, read_limit, &read_start);
    throttle(&write_bytes, write_limit, &write_start);

    ssize_t result = original_copy_file_range(fd_in, off_in, fd_out, off_out, len, flags);
    if (result > 0) {
        read_bytes += result;
        write_bytes += result;
    }
    return result;
}

int putchar_unlocked(int c) {
    if (!original_putchar_unlocked) {
        original_putchar_unlocked = dlsym(RTLD_NEXT, "putchar_unlocked");
        if (!original_putchar_unlocked) {
            fprintf(stderr, "Error loading original putchar_unlocked: %s\n", dlerror());
            exit(EXIT_FAILURE);
        }
    }
    throttle(&write_bytes, write_limit, &write_start);
    int result = original_putchar_unlocked(c);
    if (result != EOF) {
        write_bytes += 1;
    }
    return result;
}

ssize_t recv(int sockfd, void *buf, size_t len, int flags) {
    if (!original_recv) {
        original_recv = dlsym(RTLD_NEXT, "recv");
        if (!original_recv) {
            fprintf(stderr, "Error loading original recv: %s\n", dlerror());
            exit(EXIT_FAILURE);
        }
    }
    throttle(&read_bytes, read_limit, &read_start);
    ssize_t result = original_recv(sockfd, buf, len, flags);
    if (result > 0) {
        read_bytes += result;
    }
    return result;
}

ssize_t recvfrom(int sockfd, void *buf, size_t len, int flags,
                 struct sockaddr *src_addr, socklen_t *addrlen) {
    if (!original_recvfrom) {
        original_recvfrom = dlsym(RTLD_NEXT, "recvfrom");
        if (!original_recvfrom) {
            fprintf(stderr, "Error loading original recvfrom: %s\n", dlerror());
            exit(EXIT_FAILURE);
        }
    }
    throttle(&read_bytes, read_limit, &read_start);
    ssize_t result = original_recvfrom(sockfd, buf, len, flags, src_addr, addrlen);
    if (result > 0) {
        read_bytes += result;
    }
    return result;
}

ssize_t recvmsg(int sockfd, struct msghdr *msg, int flags) {
    if (!original_recvmsg) {
        original_recvmsg = dlsym(RTLD_NEXT, "recvmsg");
        if (!original_recvmsg) {
            fprintf(stderr, "Error loading original recvmsg: %s\n", dlerror());
            exit(EXIT_FAILURE);
        }
    }
    throttle(&read_bytes, read_limit, &read_start);
    ssize_t result = original_recvmsg(sockfd, msg, flags);
    if (result > 0) {
        read_bytes += result;
    }
    return result;
}

int putc(int c, FILE *stream) {
    if (!original_putc) {
        original_putc = dlsym(RTLD_NEXT, "putc");
        if (!original_putc) {
            fprintf(stderr, "Error loading original putc: %s\n", dlerror());
            exit(EXIT_FAILURE);
        }
    }
    throttle(&write_bytes, write_limit, &write_start);
    int result = original_putc(c, stream);
    if (result != EOF) {
        write_bytes += 1;
    }
    return result;
}

int putc_unlocked(int c, FILE *stream) {
    if (!original_putc_unlocked) {
        original_putc_unlocked = dlsym(RTLD_NEXT, "putc_unlocked");
        if (!original_putc_unlocked) {
            fprintf(stderr, "Error loading original putc_unlocked: %s\n", dlerror());
            exit(EXIT_FAILURE);
        }
    }
    throttle(&write_bytes, write_limit, &write_start);
    int result = original_putc_unlocked(c, stream);
    if (result != EOF) {
        write_bytes += 1;
    }
    return result;
}

size_t fread(void *ptr, size_t size, size_t nmemb, FILE *stream) {
    if (!original_fread) {
        original_fread = dlsym(RTLD_NEXT, "fread");
        if (!original_fread) {
            fprintf(stderr, "Error loading original fread: %s\n", dlerror());
            exit(EXIT_FAILURE);
        }
    }
    throttle(&read_bytes, read_limit, &read_start);
    size_t result = original_fread(ptr, size, nmemb, stream);
    if (result > 0) {
        read_bytes += result * size; 
    }
    return result;
}

size_t fread_unlocked(void *ptr, size_t size, size_t nmemb, FILE *stream) {
    if (!original_fread_unlocked) {
        original_fread_unlocked = dlsym(RTLD_NEXT, "fread_unlocked");
        if (!original_fread_unlocked) {
            fprintf(stderr, "Error loading original fread_unlocked: %s\n", dlerror());
            exit(EXIT_FAILURE);
        }
    }
    throttle(&read_bytes, read_limit, &read_start);
    size_t result = original_fread_unlocked(ptr, size, nmemb, stream);
    if (result > 0) {
        read_bytes += result * size;
    }
    return result;
}

int fgetc(FILE *stream) {
    if (!original_fgetc) {
        original_fgetc = dlsym(RTLD_NEXT, "fgetc");
        if (!original_fgetc) {
            fprintf(stderr, "Error loading original fgetc: %s\n", dlerror());
            exit(EXIT_FAILURE);
        }
    }
    throttle(&read_bytes, read_limit, &read_start);
    int result = original_fgetc(stream);
    if (result != EOF) {
        read_bytes += 1;
    }
    return result;
}

char *fgets(char *s, int size, FILE *stream) {
    if (!original_fgets) {
        original_fgets = dlsym(RTLD_NEXT, "fgets");
        if (!original_fgets) {
            fprintf(stderr, "Error loading original fgets: %s\n", dlerror());
            exit(EXIT_FAILURE);
        }
    }
    throttle(&read_bytes, read_limit, &read_start);
    char *result = original_fgets(s, size, stream);
    if (result != NULL) {
        read_bytes += strlen(result); 
    }
    return result;
}

int getc(FILE *stream) {
    if (!original_getc) {
        original_getc = dlsym(RTLD_NEXT, "getc");
        if (!original_getc) {
            fprintf(stderr, "Error loading original getc: %s\n", dlerror());
            exit(EXIT_FAILURE);
        }
    }
    throttle(&read_bytes, read_limit, &read_start);
    int result = original_getc(stream);
    if (result != EOF) {
        read_bytes += 1;
    }
    return result;
}

int getc_unlocked(FILE *stream) {
    if (!original_getc_unlocked) {
        original_getc_unlocked = dlsym(RTLD_NEXT, "getc_unlocked");
        if (!original_getc_unlocked) {
            fprintf(stderr, "Error loading original getc_unlocked: %s\n", dlerror());
            exit(EXIT_FAILURE);
        }
    }
    throttle(&read_bytes, read_limit, &read_start);
    int result = original_getc_unlocked(stream);
    if (result != EOF) {
        read_bytes += 1;
    }
    return result;
}

int scanf(const char *format, ...) {
    if (!original_scanf) {
        original_scanf = dlsym(RTLD_NEXT, "scanf");
        if (!original_scanf) {
            fprintf(stderr, "Error loading original scanf: %s\n", dlerror());
            exit(EXIT_FAILURE);
        }
    }
    throttle(&read_bytes, read_limit, &read_start); // Throttle before reading
    va_list args;
    va_start(args, format);
    int result = vscanf(format, args); 
    va_end(args);
    if (result > 0) {
        // Note: Accurate byte count for scanf is difficult; this is an estimate
        read_bytes += result;  
    }
    return result;
}

int fscanf(FILE *stream, const char *format, ...) {
    if (!original_fscanf) {
        original_fscanf = dlsym(RTLD_NEXT, "fscanf");
        if (!original_fscanf) {
            fprintf(stderr, "Error loading original fscanf: %s\n", dlerror());
            exit(EXIT_FAILURE);
        }
    }
    throttle(&read_bytes, read_limit, &read_start);
    va_list args;
    va_start(args, format);
    int result = vfscanf(stream, format, args);
    va_end(args);
    if (result > 0) {
        read_bytes += result; 
    }
    return result;
}

ssize_t send(int sockfd, const void *buf, size_t len, int flags) {
    if (!original_send) {
        original_send = dlsym(RTLD_NEXT, "send");
        if (!original_send) {
            fprintf(stderr, "Error loading original send: %s\n", dlerror());
            exit(EXIT_FAILURE);
        }
    }

    throttle(&write_bytes, write_limit, &write_start);
    ssize_t result = original_send(sockfd, buf, len, flags);
    if (result > 0) {
        write_bytes += result;
    }
    return result;
}

ssize_t sendto(int sockfd, const void *buf, size_t len, int flags,
               const struct sockaddr *dest_addr, socklen_t addrlen) {
    if (!original_sendto) {
        original_sendto = dlsym(RTLD_NEXT, "sendto");
        if (!original_sendto) {
            fprintf(stderr, "Error loading original sendto: %s\n", dlerror());
            exit(EXIT_FAILURE);
        }
    }

    throttle(&write_bytes, write_limit, &write_start);
    ssize_t result = original_sendto(sockfd, buf, len, flags, dest_addr, addrlen);
    if (result > 0) {
        write_bytes += result;
    }
    return result;
}

ssize_t sendmsg(int sockfd, const struct msghdr *msg, int flags) {
    if (!original_sendmsg) {
        original_sendmsg = dlsym(RTLD_NEXT, "sendmsg");
        if (!original_sendmsg) {
            fprintf(stderr, "Error loading original sendmsg: %s\n", dlerror());
            exit(EXIT_FAILURE);
        }
    }

    size_t total_bytes = 0;
    for (int i = 0; i < msg->msg_iovlen; i++) {
        total_bytes += msg->msg_iov[i].iov_len;
    }

    throttle(&write_bytes, write_limit, &write_start);
    ssize_t result = original_sendmsg(sockfd, msg, flags);
    if (result > 0) {
        write_bytes += total_bytes; 
    }
    return result;
}

ssize_t vmsplice(int fd, const struct iovec *iov, unsigned long nr_segs,
                 unsigned int flags) {
    if (!original_vmsplice) {
        original_vmsplice = dlsym(RTLD_NEXT, "vmsplice");
        if (!original_vmsplice) {
            fprintf(stderr, "Error loading original vmsplice: %s\n", dlerror());
            exit(EXIT_FAILURE);
        }
    }

    size_t total_bytes = 0;
    for (unsigned long i = 0; i < nr_segs; i++) {
        total_bytes += iov[i].iov_len;
    }

    throttle(&write_bytes, write_limit, &write_start); 
    ssize_t result = original_vmsplice(fd, iov, nr_segs, flags);
    if (result > 0) {
        write_bytes += total_bytes;
    }
    return result;
}

ssize_t process_vm_readv(pid_t pid, const struct iovec *local_iov, int local_iovcnt,
                         const struct iovec *remote_iov, int remote_iovcnt,
                         unsigned long flags) {
    if (!original_process_vm_readv) {
        original_process_vm_readv = dlsym(RTLD_NEXT, "process_vm_readv");
        if (!original_process_vm_readv) {
            fprintf(stderr, "Error loading original process_vm_readv: %s\n", dlerror());
            exit(EXIT_FAILURE);
        }
    }

    size_t total_bytes = 0;
    for (int i = 0; i < local_iovcnt; i++) {
        total_bytes += local_iov[i].iov_len;
    }

    throttle(&read_bytes, read_limit, &read_start);
    ssize_t result = original_process_vm_readv(pid, local_iov, local_iovcnt,
                                               remote_iov, remote_iovcnt, flags);
    if (result > 0) {
        read_bytes += total_bytes;
    }
    return result;
}

ssize_t process_vm_writev(pid_t pid, const struct iovec *local_iov, int local_iovcnt,
                          const struct iovec *remote_iov, int remote_iovcnt,
                          unsigned long flags) {
    if (!original_process_vm_writev) {
        original_process_vm_writev = dlsym(RTLD_NEXT, "process_vm_writev");
        if (!original_process_vm_writev) {
            fprintf(stderr, "Error loading original process_vm_writev: %s\n", dlerror());
            exit(EXIT_FAILURE);
        }
    }

    size_t total_bytes = 0;
    for (int i = 0; i < local_iovcnt; i++) {
        total_bytes += local_iov[i].iov_len;
    }

    throttle(&write_bytes, write_limit, &write_start);
    ssize_t result = original_process_vm_writev(pid, local_iov, local_iovcnt,
                                                remote_iov, remote_iovcnt, flags);
    if (result > 0) {
        write_bytes += total_bytes;
    }
    return result;
}

__attribute__((constructor))
void init() {
    clock_gettime(CLOCK_MONOTONIC, &read_start);
    clock_gettime(CLOCK_MONOTONIC, &write_start);

    // Initialize limits from environment variables (bytes/s)
    char *read_env = getenv("IOTHROTTLE_READ");
    char *write_env = getenv("IOTHROTTLE_WRITE");
    char *debug_env = getenv("IOTHROTTLE_DEBUG");

    if (debug_env && atoi(debug_env) != 0) {
        debug_mode = 1;
    }

    if (debug_mode) fprintf(stderr, "init called: default read_limit=%zu, write_limit=%zu\n", read_limit, write_limit);

    if (read_env) {
        if (debug_mode) fprintf(stderr, "IOTHROTTLE_READ=%s\n", read_env);
        read_limit = strtoull(read_env, NULL, 10);
    } else {
        if (debug_mode) fprintf(stderr, "IOTHROTTLE_READ not set\n");
    }

    if (write_env) {
        if (debug_mode) fprintf(stderr, "IOTHROTTLE_WRITE=%s\n", write_env);
        write_limit = strtoull(write_env, NULL, 10);
    } else {
        if (debug_mode) fprintf(stderr, "IOTHROTTLE_WRITE not set\n");
    }

    if (debug_mode) fprintf(stderr, "Final read_limit=%zu, write_limit=%zu\n", read_limit, write_limit);
}