#include <unistd.h>
#include <limits.h>
#include <string.h>
#include <errno.h>
#include <stdio.h>
#include <stdlib.h>
#include <fcntl.h>
#include <sys/wait.h>

#define MIN_MAX_ARG 4096
#define MIN_MAX_ENV 4096
#define STEPPING 1024

#define TEST_PROG "/bin/echo"

static int verbose = 0;

void fill_data(char* location, int amount) {
    char x = 'a';
    while(amount-- > 0) {
        if(x > 'z') x = 'a';
        location[amount] = x++;
    }
}

char ** alloc_arr(int sz, int for_env) {
    char ** arr = malloc(sizeof(char*) * (for_env == 0 ? 3 : 2));
    int data_index = 0;

    if(!for_env) {
        arr[0] = strdup(TEST_PROG);
        sz -= strlen(arr[0]) + 1;
        data_index++;
    }

    arr[data_index] = malloc(sz + 1);

    if(!arr[data_index]) {
        perror("malloc");
        exit(1);
    }

    fill_data(arr[data_index], sz);
    arr[data_index][sz] = 0;
    
    data_index++;
    arr[data_index] = NULL;

    return arr;
}

void free_arr(char** arr) {
    int i = 0;

    while(arr[i])
        free(arr[i++]);

    free(arr);
}

int calc_env_sz(char** env) {
    int acc = 0;

    while(env && *env) {
        acc += strlen(*env) + 1;
        ++env;
    }

    return acc;
}

int test_exec_env(int start_sz, int cmd_size) {
    char ** arg = alloc_arr(cmd_size, 0);

    while(1) {
        char ** env = alloc_arr(start_sz, 1);
        if(verbose)
            printf("trying env with sz: %d\n", start_sz);

        if(try_fork_exec(arg, env) == 0) {
            free_arr(env);
            free_arr(arg);
            return start_sz;
        }

        free_arr(env);
        start_sz -= STEPPING;

        if(start_sz < 0) {
            printf("oups - env size negative!\n");
            exit(1);
        }
    }
}

int test_exec_arg(int start_sz, int env_sz, int once) {
    char ** env = alloc_arr(env_sz, 1);

    while(1) {
        char ** cmd = alloc_arr(start_sz, 0);
        if(verbose)
            printf("trying cmd with sz: %d\n", start_sz);

        if(try_fork_exec(cmd, env) == 0) {
            free_arr(cmd);
            free_arr(env);
            return start_sz;
        }

        free_arr(cmd);
        start_sz -= STEPPING;

        /* didn't work out */
        if(once) return -1;

        if(start_sz < 0) {
            printf("oups - arg size negative!\n");
            exit(1);
        }
    }
}

int try_fork_exec(char ** argv, char** env) {
    int pid = fork();
    switch(pid) {
    case 0: {
        /* child */
        int new_so = open("/dev/null", O_WRONLY);
        if(new_so == -1) {
            perror("open");
            exit(2);
        }
        if(dup2(new_so, fileno(stdout)) == -1) {
            perror("dup2");
            exit(2);
        }
        
        execve(argv[0], argv, env);

        if(verbose)
            perror("exec");
        exit(3);
        break;
    }
    case -1:
        /* error */
        perror("fork");
        exit(1);
        break;
    default: {
        /* parent */
        int status;
        waitpid(pid, &status, 0);

        /* exited */
        if(!WIFEXITED(status)) {
            perror("waitpid");
            exit(1);
        }
        
        return WEXITSTATUS(status);
    }
    }
}

int main(int argc, char** argv) {
    int max_arg = -1;
    int max_env = -1;

#ifdef _SC_ARG_MAX
    printf("using sysconf...\n");
    max_arg = sysconf(_SC_ARG_MAX);
#endif

#ifdef ARG_MAX
    if(max_arg == -1) {
      printf("using ARG_MAX...\n");
      max_arg = ARG_MAX;
    }
#endif

    if(max_arg == -1) {
      printf("using fallback...\n");
      max_arg = 4096;
    }

    printf("max_arg: %d\n", max_arg);

#if 1
    extern char** environ;
    printf("env_sz : %d\n", calc_env_sz(environ));
#endif
    max_env = max_arg;

    printf("testing...\n");

    max_env = test_exec_env(max_env, MIN_MAX_ARG);
    printf("max_env with min_arg (%d): %d\n", MIN_MAX_ARG, max_env);

    max_arg = test_exec_arg(max_arg, MIN_MAX_ENV, 0);
    printf("max_arg with min_env (%d): %d\n", MIN_MAX_ENV, max_arg);

    printf("max_env with max_arg (%d): %d\n", max_arg, test_exec_env(max_env, max_arg));
    printf("max_arg with max_env (%d): %d\n", max_env, test_exec_arg(max_arg, max_env, 0));

    {
        int tmp;

        tmp = test_exec_arg(max_arg, max_env, 1);
        printf("is max_env & max_arg together ok? %s\n",  tmp == -1 ? "NOT OK" : "ok");

        if(tmp == -1) {
            tmp = test_exec_arg((max_arg / 2) - 1, (max_env / 2) - 1, 1);
            printf("is max_env / 2 & max_arg / 2 together ok? %s\n", tmp == -1 ? "NOT OK" : "ok");

            if(tmp == -1) {
                tmp = test_exec_arg((max_arg / 4) - 1, (max_env / 4) - 1, 1);
                printf("is max_env / 4 & max_arg / 4 together ok? %s\n", tmp == -1 ? "NOT OK" : "ok");

                if(tmp == -1) {
                    printf("failed to find reasonable limits! giving up...\n");
                    exit(1);
                } else {
                    max_arg = (max_arg / 4) - 1;
                    max_env = (max_env / 4) - 1;
                }
            } else {
                max_arg = (max_arg / 2) - 1;
                max_env = (max_env / 2) - 1;
            }
        }
    }

    printf("---------------------------------------\n");
    printf("limits: arg: %d, env: %d\n", max_arg, max_env);

    return 0;
}