#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>
#include <signal.h>

using namespace std;

vector<string> functions = {"exit", "cd", "sub", "listsubs", "set", "get", "run"};

vector<pair<string, vector<string>>> substitutions;

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;
}

vector<string> getPath() {
    string pathstr = getenv("PATH");
    if (pathstr == "") {
        log("WARNING! SPACE CANNOT FIND THE PATH! Please ensure your path is set as an environment variable in your Spacefile, or one of Space's parents has set the PATH environment variable. For now, enjoy the default path :)");
    } else {
        vector<string> path;
        string buffer;
        for (int i = 0; i < pathstr.size(); i++) {
            if (pathstr[i] == ':') {
                path.push_back(buffer);
                buffer = "";
            } else {
                buffer += pathstr[i];
            }
        }
        path.push_back(buffer);
        return path;
    }
    return {"/bin", "/usr/bin", "/usr/local/bin"}; 
}

string findExecutable(string executable) {
    for (int i = 0; i < functions.size(); i++) {
        if (functions[i] == executable) {
            return executable;
        }
    }
    if (filesystem::exists(executable)) {
        if (executable[0] != '/') {
            return "./" + executable;
        }
        return executable;
    }
    string output;
    vector<string> path = getPath();
    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) {
    if (args[0] == "exit") {
        exit(0);
    } else if (args[0] == "cd") {
        if (args.size() == 1) {
            log("cd requires an argument");
            return 1;
        } else {
            filesystem::current_path(args[1]);
        }
    } else if (args[0] == "set") {
        if (args.size() < 3) {
            log("set requires 2 arguments");
            return 1;
        }
        string envStr = (args[1] + "=" + args[2]);
        char* cEnvStr = strdup(envStr.c_str());
        putenv(cEnvStr);
    } else if (args[0] == "get") {
        if (args.size() < 2) {
            log("get requires an argument");
            return 1;
        }
        char* value = getenv(args[1].c_str());
        if (value == nullptr) {
            log("Environment variable " + args[1] + " not found");
            return 1;
        }
        cout << value << endl;
    } else if (args[0] == "sub") {
        if (args.size() < 3) {
            log("substitute requires 2 arguments");
            return 1;
        }
        vector<string> subArgs = tokeniseSubstitutions(args[2]);
        substitutions.push_back({args[1], subArgs});
    } else if (args[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;
        }
    } else {
        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; 
    }
    return 0;
}

void runScript(string script) {
    string scriptLine;
    ifstream scriptFile(script);
    while (getline(scriptFile, scriptLine)) {
        vector<string> tokens = tokenise(scriptLine);
        if (tokens.empty()) continue;
        int returnCode = runProcess(tokens);
        if (returnCode != 0) {
            log("Script error:\nCommand " + tokens[0] + " failed with exit code " + to_string(returnCode));
        }
    }
    scriptFile.close();
}

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");
    }
}

void handler(int signum);

int main(int argc, char *argv[]) {
    signal(SIGINT, handler);
    if (argc > 1) {
        runScript(argv[2]);
    }
    doStartup();
    bool continueLoop = true;
    while (continueLoop) {
        string input = prompt();
        if (input == "") {
            continue;
        }
        vector<string> tokens = tokenise(input);
        auto path = filesystem::current_path();
        if (tokens.empty()) {
            continue;
        }
        int returnCode = runProcess(tokens);
        if (returnCode != 0) {
            log("Process failed with error code " + to_string(returnCode));
        }
    }
    return 0;
}

void handler(int signum) {
    main(1, nullptr);
}