#define _GNU_SOURCE

#include <stdlib.h>
#include <stdio.h>
#include <string.h>
#include <unistd.h>
#include <sys/wait.h>
#include <sys/utsname.h>
#include <errno.h>
#include <termios.h>
#include <fcntl.h>
#include <sys/stat.h>

#include "ansi.h"


#define HISTORY_FILE_NAME "/history.log"
#define LOCAL_SHARE_DIRNAME "/.local/share/msh"

#define DEFAULTLINECAPACITY     256
#define DEFAULTTOKENSCAPACITY   64
#define DEFAULTHISTORYCAPACITY  128


int cd(int argc, char** argv) {
    if (argc != 1) {
        fprintf(stderr, "Usage: cd <directory>\n");
        return 1;
    }

    return chdir(argv[0]);
}


void pretty_path_home(char** path) {
    char* home = secure_getenv("HOME");
    int home_len = strlen(home);

    if (strncmp(home, *path, home_len) == 0) {
        (*path)[home_len-1] = '~';
        *path = *path + (home_len - 1);
    }
}


int read_ansi() {
    switch (getchar()) {
        case '[':
            switch(getchar()) {
                case 'A':
                    return ANSI_CURSOR_UP;
                case 'B':
                    return ANSI_CURSOR_DOWN;
                case 'C':
                    return ANSI_CURSOR_RIGHT;
                case 'D':
                    return ANSI_CURSOR_LEFT;
                case '1':
                    switch(getchar()) {
                        case ';':
                            switch(getchar()) {
                                // CTRL
                                case '5':
                                // ALT
                                case '3':
                                    switch (getchar()) {
                                        case 'C':
                                            return ANSI_CURSOR_RIGHT_WORD;
                                        case 'D':
                                            return ANSI_CURSOR_LEFT_WORD;
                                    }
                                    break;
                            }
                            break;
                    }
                    break;
            }
            break;
        case 127:
        case 8:
            return ANSI_DELETE_WORD;
        default:
            return ANSI_UNKNOWN;
    }

    return ANSI_UNKNOWN;
}


// Find the length of the prefix until the last occurrence of a rejected character
size_t strnrcspn(char* str, size_t n, char* reject) {
    size_t i;
    int seenaccept = 0;

    for (i = n - 1; i > 0; i--) {
        // If this is NOT a reject character
        if(strchr(reject, str[i]) == NULL) {
            seenaccept = 1;
            continue;
        } else {
            if (seenaccept == 1) {
                break;
            }
        }
    }

    return i;
}

// Find the length of the prefix until the first occurrence of a rejected character after an offset
size_t strnocspn(char* str, size_t n, size_t o, char* reject) {
    size_t i;
    int seenaccept = 0;

    for (i = o; i < n; i++) {
        // If this is NOT a reject character
        if(strchr(reject, str[i]) == NULL) {
            seenaccept = 1;
            continue;
        } else {
            if (seenaccept == 1) {
                break;
            }
        }
    }

    return i;
}


char* status_string(int status) {
    switch (status) {
        case 0:
            return "";
    }

    return "❗";
}


int main(int argc, char** argv) {
    int retval = 0;
    char*   line = NULL;
    size_t  linecapacity = 0;
    char    prompt[DEFAULTLINECAPACITY] = {};
    size_t  linelen = 0;
    char*   tokenwork = NULL;
    size_t  tokenworkcapacity = 0;
    char**  tokens = NULL;
    size_t  tokenscapacity = 0;
    size_t  ntokens = 0;
    char**  history = NULL;
    size_t  historycapacity = 0;
    size_t  historyi = 0;
    size_t  historyoffset = 0;
    char*   oldline = NULL;
    

    int pid;
    int laststatus = 0;

    int wstatus;
    char pwd[DEFAULTLINECAPACITY] = {};
    char* pwd_disp;

    char* user;
    char* host;

    struct utsname uts = {};
    char* sharedir = NULL;
    char* historyfilepath = NULL;
    FILE* historyfile = NULL;

    user = secure_getenv("USER");

    int result = uname(&uts);
    if (result != 0) {
        fprintf(stderr, "Error: %s\n", strerror(errno));
        retval = errno;
        goto cleanup;
    }

    host = uts.nodename;


    struct termios term, orig_term;
    tcgetattr(fileno(stdin), &term);
    orig_term = term;

    // disable echo and icanon (don't wait for newline)
    term.c_lflag &= ~(ECHO | ICANON);

    char* homedir = secure_getenv("HOME");

    if (homedir == NULL) {
        fprintf(stderr, "Error: $HOME is not specified\n");
        retval = 1;
        goto cleanup;
    }

    sharedir = malloc(strlen(homedir) + strlen(LOCAL_SHARE_DIRNAME) + 1);
    if (sharedir == NULL) {
        retval = 1;
        goto cleanup;
    }
    sharedir = strcpy(sharedir, homedir);
    if (sharedir == NULL) {
        retval = 1;
        goto cleanup;
    }
    sharedir = strcat(sharedir, LOCAL_SHARE_DIRNAME);
    if (sharedir == NULL) {
        retval = 1;
        goto cleanup;
    }

    result = mkdir(sharedir, S_IRWXU);
    if (result != 0 && errno != EEXIST) {
        fprintf(stderr, "Error: %s\n", strerror(errno));
        retval = 1;
        goto cleanup;
    }


    historyfilepath = malloc(strlen(sharedir) + strlen(HISTORY_FILE_NAME) + 1);
    if (historyfilepath == NULL) {
        retval = 1;
        goto cleanup;
    }
    historyfilepath = strcpy(historyfilepath, sharedir);
    if (historyfilepath == NULL) {
        retval = 1;
        goto cleanup;
    }
    historyfilepath = strcat(historyfilepath, HISTORY_FILE_NAME);
    if (historyfilepath == NULL) {
        retval = 1;
        goto cleanup;
    }

    int historyfilefd = open(historyfilepath, O_CREAT | O_APPEND | O_WRONLY, S_IRUSR | S_IWUSR);
    if (historyfilefd == -1) {
        fprintf(stderr, "Error: could not open history file: %s\n", strerror(errno));
        retval = 1;
        goto cleanup;
    }

    historyfile = fdopen(historyfilefd, "w");
    if (historyfile == NULL) {
        fprintf(stderr, "Error: could not open history file fd: %s\n", strerror(errno));
        retval = 1;
        goto cleanup;
    }


    // Initialize memory
    tokens = calloc(DEFAULTTOKENSCAPACITY, sizeof(char*));
    if (tokens == NULL) {
        retval = 1;
        goto cleanup;
    }
    memset(tokens, 0, DEFAULTTOKENSCAPACITY * sizeof(char*));
    tokenscapacity = DEFAULTTOKENSCAPACITY;

    tokenwork = malloc(DEFAULTLINECAPACITY);
    if (tokenwork == NULL) {
        retval = 1;
        goto cleanup;
    }
    tokenworkcapacity = DEFAULTLINECAPACITY;

    historycapacity = DEFAULTHISTORYCAPACITY;
    history = calloc(historycapacity, sizeof(char*));
    if (history == NULL) {
        retval = 1;
        goto cleanup;
    }
    memset(history, 0, historycapacity);

    // Main loop
    while(1) {
        char* result = getcwd(pwd, DEFAULTLINECAPACITY - 1);
        if (result == NULL) {
            retval = 1;
            goto cleanup;
        }

        pwd_disp = pwd;

        pretty_path_home(&pwd_disp);


        snprintf(
            prompt,
            DEFAULTLINECAPACITY,

            ANSI_BG_BLUE ANSI_FG_BLACK
            "%s"                                  // status
            " %s"                                 // user
            "@"
            ANSI_FG_BLACK
            "%s "                                 // host
            ANSI_COLOR_RESET ANSI_FG_BLUE
            " "
            ANSI_FG_BLUE
            "%s"                                  // directory
            ANSI_COLOR_RESET
            "$ "
            ANSI_COLOR_RESET,

            status_string(laststatus), user, host, pwd_disp
        );
        printf("%s", prompt);


        // prepare for input
        linelen = 0;
        int readc;
        int cursor = 0;

        tcsetattr(fileno(stdin), 0, &term);

        line = malloc(DEFAULTLINECAPACITY);
        linecapacity = DEFAULTLINECAPACITY;

        memset(line, 0, linecapacity);

        historyoffset = 0;

        while ((readc = getchar()) != '\n') {
            if (readc == EOF) {
                printf("^D\n");
                retval = 0;
                goto cleanup;
            }

            if (readc == '\x1b') {
                int ansi = read_ansi();
                if (ansi == ANSI_CURSOR_LEFT) {
                    if (cursor > 0) {
                        cursor--;
                    }
                } else if (ansi == ANSI_CURSOR_RIGHT) {
                    if (cursor < linelen) {
                        cursor++;
                    }

                } else if (ansi == ANSI_CURSOR_UP || ansi == ANSI_CURSOR_DOWN) {
                    int load_from_hist = 1;

                    if (ansi == ANSI_CURSOR_UP && historyoffset < historyi) {
                        historyoffset++;

                        if (historyoffset == 1) {
                            if (oldline != NULL) {
                                free(oldline);
                            }

                            oldline = line;
                            line = NULL;
                        }
                    } else if (ansi == ANSI_CURSOR_DOWN && historyoffset > 0) {
                        historyoffset--;

                        // We just went back to the last line, load from oldline
                        if (historyoffset == 0) {
                            free(line);
                            line = oldline;
                            linelen = strlen(line);
                            linecapacity = linelen + 1;
                            cursor = linelen;

                            oldline = NULL;

                            // Don't look up from history
                            load_from_hist = 0;
                        }
                    } else {
                        // We didn't change our line (already at top or bottom)
                        load_from_hist = 0;
                    }


                    if (load_from_hist == 1) {
                        char*  historyline = history[historyi - (historyoffset)];
                        size_t historylinelen = strlen(historyline);

                        if (line != NULL) {
                            free(line);
                        }

                        line = malloc(historylinelen + 1);
                        if (line == NULL) {
                            retval = 1;
                            goto cleanup;
                        }
                        line = strcpy(line, historyline);
                        if (line == NULL) {
                            retval = 1;
                            goto cleanup;
                        }

                        cursor = historylinelen;
                        linecapacity = historylinelen + 1;
                        linelen = historylinelen;
                    }
                } else if (ansi == ANSI_DELETE_WORD) {
                    if (cursor > 0) {
                        // 0123 56
                        size_t index = strnrcspn(line, cursor, " \t");
                        memmove(line + index, line + cursor, linecapacity - cursor);
                        linelen -= cursor - index;
                        cursor = index;
                    }
                } else if (ansi == ANSI_CURSOR_LEFT_WORD) {
                    size_t index = strnrcspn(line, cursor, " \t");
                    cursor = index;
                } else if (ansi == ANSI_CURSOR_RIGHT_WORD) {
                    size_t index = strnocspn(line, linelen, cursor, " \t");
                    cursor = index;
                }
            } else if (readc == 127) {
                // backspace
                if (cursor > 0) {
                    cursor--;
                    linelen--;
                }

                if (cursor == linelen) {
                    line[cursor] = 0;
                } else {
                    memmove(line + cursor, line + cursor + 1, linecapacity - (cursor));
                }
            } else {
                // We are typing
                if (cursor != linelen) {
                    memmove(line + cursor + 1, line + cursor, linecapacity - cursor - 1);
                }

                line[cursor] = readc;
                cursor++;
                linelen++;

                // Are we out of capacity? Double line length
                while (linelen >= linecapacity - 1) {
                    linecapacity *= 2;
                    line = realloc(line, linecapacity);
                    if (line == NULL) {
                        fprintf(stderr, "Error: Failed to re-allocated line buffer");
                        retval = 1;
                        goto cleanup;
                    }

                    memset(line + linecapacity / 2, 0, linecapacity / 2);
                }
            }

            // Manual echo
            printf(ANSI_ESC "[2K" ANSI_ESC "[0G" "%s%s", prompt, line);
            for (int j = 0; j < linelen - cursor; j++) {
                printf(ANSI_ESC "[D");
            }
            fflush(stdout);
        }

        // oldline is only valid while reading
        if (oldline != NULL) {
            free(oldline);
            oldline = NULL;
        }

        putc('\n', stdout);

        tcsetattr(fileno(stdin), 0, &orig_term);

        line[linelen] = 0;

        if (linelen == 0) {
            free(line);
            line = NULL;
            continue;
        }

        fprintf(historyfile, "%s\n", line);
        fflush(historyfile);
        
        while (tokenworkcapacity + 1 < linelen) {
            tokenworkcapacity *= 2;
            tokenwork = realloc(tokenwork, tokenworkcapacity);
        }

        memset(tokenwork, 0, tokenworkcapacity);
        tokenwork = strncpy(tokenwork, line, linelen);
        if (tokenwork == NULL) {
            retval = 1;
            goto cleanup;
        }


        while (historyi >= historycapacity) {
            historycapacity *= 2;
            history = reallocarray(history, historycapacity, sizeof(char*));

            if (history == NULL) {
                retval = 1;
                goto cleanup;
            }
            memset(history + historyi, 0, (historycapacity - historyi) * sizeof(char*));
        }
        // change ownership of line to history
        history[historyi++] = line;
        line = NULL;

        char* tokensarg = tokenwork;
        ntokens = 0;

        // Split into tokens
        for (;; ntokens++) {
            char* token = strtok(tokensarg, " \t");
            if (token == NULL) {
                break;
            }

            tokens[ntokens] = token;

            while (ntokens >= tokenscapacity - 1) {
                tokenscapacity *= 2;
                tokens = reallocarray(tokens, tokenscapacity, sizeof(char*));
            }

            // required to be NULL for second strtok invocation
            if (tokensarg != NULL) {
                tokensarg = NULL;
            }
        }

        tokens[ntokens] = NULL;
        char** args = &tokens[1];

        if (ntokens == 0) {
            continue;
        }

        // Shell builtins
        if (strcmp("exit", tokens[0]) == 0) {
            printf("Goodbye!\n");
            retval = 0;
            goto cleanup;
        }

        if (strcmp("cd", tokens[0]) == 0) {
            laststatus = cd(ntokens - 1, args);
            continue;
        }

        if (strcmp("history", tokens[0]) == 0) {
            for (int i = 0; i < historycapacity; i++) {
                if (history[i] == NULL) {
                    break;
                }

                printf("%.3d: %s\n", i, history[i]);
            }

            laststatus = 0;
            continue;
        }

        // Execute command
        pid = fork();

        if(pid == 0) {
            retval = execvpe(tokens[0], tokens, environ);
            goto cleanup;
        } else {
            waitpid(pid, &wstatus, 0);

            laststatus = WEXITSTATUS(wstatus);
        }


    }


cleanup:
    if (sharedir != NULL) {
        free(sharedir);
    }
    if (historyfilepath != NULL) {
        free(historyfilepath);
    }

    if (historyfile != NULL) {
        fflush(historyfile);
        fclose(historyfile);
    }

    if (line != NULL) {
        free(line);
    }

    if (oldline != NULL) {
        free(oldline);
    }

    if (tokenwork != NULL) {
        free(tokenwork);
    }

    if (tokens != NULL) {
        free(tokens);
    }

    if (history != NULL) {
        for (int i = 0; i < historycapacity; i++) {
            if (history[i] != NULL) {
                free(history[i]);
            }
        }

        free(history);
    }


    tcsetattr(fileno(stdin), 0, &orig_term);
    return retval;
}