#include <iostream>
#include <string>
#include <vector>
#include <ctime>
#include <filesystem>
#include <unistd.h>
#include <sys/wait.h>
#include <sys/types.h>
#include <pwd.h>
#include <cstring>
#include <limits.h>
#include <fstream>

using namespace std;

vector<string> builtInFunctions = {"exit", "cd", "sub"};

vector<pair<string, vector<string>>> substitutions = {{"ls", {"ls", "--color"}}};

void log(string input) {
    cout << input << endl;
}

string getTime() {
    time_t rawtime;
    struct tm * timeinfo;
    char buffer[40];

    time(&rawtime);
    timeinfo = localtime(&rawtime);

    strftime(buffer, sizeof(buffer),"%H:%M:%S", timeinfo);
    string str(buffer);
    return str;
}

string prompt() {
    string input;
    char hostname[HOST_NAME_MAX];
    gethostname(hostname, HOST_NAME_MAX);
    cout << "\x1B[0m{" << getTime() << "}\x1B[32m " << getenv("USER") << "@" << hostname << "\x1B[96m " << string(filesystem::current_path()) << endl << "\x1B[0m> ";
    getline(cin, input);
    return input;
}

string findExecutable(string executable) {
    for (int i = 0; i < builtInFunctions.size(); i++) {
        if (builtInFunctions[i] == executable) {
            return executable;
        }
    }
    if (filesystem::exists(executable)) {
        if (executable[0] != '/') {
            return "./" + executable;
        }
        return executable;
    }
    string output;
    vector<string> path = {"/bin", "/usr/bin", "/usr/local/bin"};
    for (int i = 0; i < path.size(); i++) {
        output = path[i] + "/" + executable;
        if (filesystem::exists(output)) {
            return output;
        } else {
            output.clear();
        }
    }
    return output;
}

string getHomeDir() {
    struct passwd* pw = getpwuid(getuid());
    if (pw && pw->pw_dir) return string(pw->pw_dir);
    return "";
}

vector<string> tokenise(string input) {
    vector<string> output;
    string currentArg;
    bool isString = false;
    for (int i = 0; i < input.size(); i++) {
        if (input[i] == '"') {
            isString = !isString;
        } else if (input[i] == ' ' && !currentArg.empty() && !isString) {
            output.push_back(currentArg);
            currentArg.clear();
        } else if (input[i] == '~') {
            currentArg += getHomeDir();
        } else {
            currentArg += input[i];
        }
    }
    if (!currentArg.empty()) {
        output.push_back(currentArg);
    }
    for (int i = 0; i < substitutions.size(); i++) {
        for (int j = 0; j < output.size(); j++) {
            if (output[j] == substitutions[i].first) {
                output.erase(output.begin() + j);
                int substitutionSize = substitutions[i].second.size();
                for (int k = 0; k < substitutionSize; k++) {
                    output.insert(output.begin() + j, substitutions[i].second[k]);
                    j++;
                }
                j -= substitutionSize;
                break;
            }
        }
    }
    string executableLocation = findExecutable(output[0]);
    if (!executableLocation.empty()) {
        output[0] = executableLocation;
    } else {
        log("Couldn't find an executable in the path");
    }
    return output;
}

vector<string> tokeniseSubstitutions(string input) {
    vector<string> output;
    string currentArg;
    bool isString = false;
    for (int i = 0; i < input.size(); i++) {
        if (input[i] == '"') {
            isString = !isString;
        } else if (input[i] == ' ' && !isString) {
            if (!currentArg.empty()) {
                output.push_back(currentArg);
                currentArg.clear();
            }
        } else {
            currentArg += input[i];
        }
    }
    if (!currentArg.empty()) {
        output.push_back(currentArg);
    }
    return output;
}

int runProcess(vector<string> args) {
    pid_t pid = fork();
    if (pid == -1) {
        cerr << "Fork failed!" << endl;
        return -1;
    }
    if (pid == 0) {
        vector<char*> cargs;
        cargs.push_back(const_cast<char*>(args[0].c_str()));        
        for (int i = 1; i < args.size(); i++) {
            cargs.push_back(const_cast<char*>(args[i].c_str()));        
        }
        cargs.push_back(nullptr);
        execvp(args[0].c_str(), cargs.data());
        
        cerr << "Execution failed with error " << strerror(errno) << endl;
        _exit(1);
    }
    int status;
    waitpid(pid, &status, 0);

    if (WIFEXITED(status)) {
        return WEXITSTATUS(status);
    }
    return -1;
}

void doStartup() {
    struct passwd* pw = getpwuid(getuid());
    const char* homedir = pw->pw_dir;
    string configLocation = string(homedir) + "/.config/space/Spacefile";
    if (filesystem::exists(configLocation)) {
        string config;
        ifstream configFile(configLocation);
        while (getline(configFile, config)) {
            vector<string> tokens = tokenise(config);
            if (tokens.empty()) continue;
            int returnCode = runProcess(tokens);
            if (returnCode != 0) {
                log("Config file error:\nCommand " + tokens[0] + " failed with exit code " + to_string(returnCode));
            }
        }
        configFile.close();
    } else {
        log("Spacefile does not exist");
    }
}

int main(int argc, char *argv[]) {
    doStartup();
    bool continueLoop = true;
    while (continueLoop) {
        string input = prompt();
        if (input == "") {
            continue;
        }
        vector<string> tokens = tokenise(input);
        auto path = filesystem::current_path();
        if (tokens[0] == "exit") {
            exit(0);
        } else if (tokens[0] == "cd") {
            if (tokens.size() == 1) {
                log("cd requires an argument");
            } else {
                filesystem::current_path(tokens[1]);
            }
            continue;
        } else if (tokens[0] == "sub") {
            if (tokens.size() < 3) {
                log("substitute requires 2 arguments");
                continue;
            }
            vector<string> args = tokeniseSubstitutions(tokens[2]);
            substitutions.push_back({tokens[1], args});
            continue;
        } else if (tokens[0] == "listsubs") {
            log("Substitutions:");
            for (int i = 0; i < substitutions.size(); i++) {
                string substring;
                for (int j = 0; j < substitutions[i].second.size(); j++) {
                    substring += substitutions[i].second[j] + " ";
                }
                cout << substitutions[i].first << " -> " << substring << endl;
            }
        }
        if (tokens.empty()) {
            continue;
        }
        int returnCode = runProcess(tokens);
        if (returnCode != 0) {
            log("Process failed with error code " + to_string(returnCode));
        }
    }
    return 0;
}