#include "Shader.h"

namespace NB{

FileLocation get_file_path(std::string name) {
    const std::vector<char> allowed_special_chars = {
        '!', '#', '$', '%', '&', '\'', '(', ')', '+', ',', '-', ';', '=', '@', '[', ']', '^', '_', '`', '~' 
    };
    std::vector<std::string> path = {""};
    for (const char c : name) {
        if (c == '/' || c == '\\') {
            path.push_back("");
        } else if (std::isalnum(c)) {
            path.back() += c;
        } else {
            for (const char sc : allowed_special_chars) {
                if (c == sc) {
                    path.back() += c;
                }
            }
            throw std::runtime_error("Not valid filepath.");
        }
    }
    FileLocation ret;
    for (const auto& tk : path) {
        if (tk == path.back()) {
            ret.second = tk;
        } else {
            ret.first += '/' + tk;
        }
    }
    return ret;
}

// ShaderPreprocessorError class
ShaderPreprocessorError::ShaderPreprocessorError(
    const std::string& msg,
    const std::string& file,
    int line
) : code(Codes::UNDEFINED), std::runtime_error(formatDebugString(msg, file, line)) {}

ShaderPreprocessorError::ShaderPreprocessorError(
    Codes err_code,
    const std::string& arg,
    const std::string& file,
    int line
) : code(err_code), std::runtime_error(formatDebugString(errorCodeParser(err_code, arg), file, line)) {}

std::string ShaderPreprocessorError::errorCodeParser(Codes err_code, const std::string& arg) {
    switch(err_code) {
        case Codes::FILE_NOT_FOUND:
            return "File '" + arg + "' not found.";
        case Codes::BUILTIN_NOT_FOUND:
            return "Built-in '" + arg + "' not found.";
        case Codes::CUSTOM:
        case Codes::UNDEFINED:
            return arg;
        case Codes::NONE:
        default:
            return "";
    }
}

// ShaderError class
ShaderError::ShaderError(
    const std::string& msg,
    const char* shad,
    const std::string& file,
    int line
) : std::runtime_error(formatString(msg, shad, file, line)) {}

std::string ShaderError::formatString(
    const std::string& msg,
    const char* shad,
    const std::string& file,
    int line
) {
    std::string ret = "";
    if (file != "") {
        ret += "In file " + file;
        if (line >= 0) {
            ret += " at line " + std::to_string(line);
        }
        ret += ":\n\t";
    }
    ret += msg;
    if (shad != nullptr) {
        ret += " with shader error: " + std::string(shad);
    }
    return ret;
}

void preprocessor_include(const ShaderPreprocessor& proc, ShaderUnit& shad, unsigned int linenum, const std::string& line) {
    typedef ShaderPreprocessor::Codes Codes;
    std::stringstream lstream(line);
    std::vector<std::string> tokens;
    std::string tk;
    while (std::getline(lstream, tk, ' ')) {
        tokens.push_back(tk);
    }
    // Error? When #include isn't at beginning of line
    if (tokens[0] != "#include") { return; }
    std::stringstream include_code;
    std::regex dir_pattern("\".*\"");
    std::cmatch matches;
    std::stringstream raw;
    std::string filename;
    std::regex bin_pattern("<.*>");
    if (std::regex_match(tokens[1].c_str(), matches, dir_pattern)) {
        filename = matches[0];
        filename = filename.substr(1, filename.size()-2);
        try {
            raw = proc.loadFromRelative(filename, shad.file.loc.first);
            proc.process(raw, shad);
            std::cout << "b1\n";
        } catch (ShaderPreprocessorError e) {
            if (e.code == Codes::FILE_NOT_FOUND) {
                try {
                    raw = proc.loadFromDirectory(filename);
                    proc.process(raw, shad);
                } catch (ShaderPreprocessorError f) {
                    if (f.code == Codes::FILE_NOT_FOUND) {
                        raw = proc.loadFromBuiltIn(filename);
                        proc.process(raw, shad);
                    } else {
                        throw f;
                    }
                }
            } else {
                throw e;
            }
        }
    } else if (std::regex_match(tokens[1].c_str(), matches, bin_pattern)) {
        filename = matches[0];
        filename = filename.substr(1, filename.size()-2);
        try {
            raw = proc.loadFromDirectory(filename);
            proc.process(raw, shad);
        } catch (ShaderPreprocessorError e) {
            if (e.code == Codes::FILE_NOT_FOUND) {
                try {
                    raw = proc.loadFromBuiltIn(filename);
                    proc.process(raw, shad);
                } catch (ShaderPreprocessorError f) {
                    if (f.code == Codes::BUILTIN_NOT_FOUND) {
                        raw = proc.loadFromRelative(filename, shad.file.loc.first);
                        proc.process(raw, shad);
                    } else {
                        throw f;
                    }
                }
            } else {
                throw e;
            }
        }
    }

}

void preprocessor_define(const ShaderPreprocessor& proc, ShaderUnit& shad, unsigned int linenum, const std::string& line) {
    std::stringstream lstream(line);
    std::vector<std::string> tokens;
    std::string tk;
    while (std::getline(lstream, tk, ' ')) {
        tokens.push_back(tk);
    }
    if (tokens[0] != "#define") { return; }
    shad.defines[tokens[1]] = (tokens.size() > 2) ? tokens[2] : "";
    shad.preprocSource << line << "\n";
}

void preprocessor_version(const ShaderPreprocessor& proc, ShaderUnit& shad, unsigned int linenum, const std::string& line) {
    shad.preprocSource << line << "\n" ;
}

// ShaderPreprocessor class
ShaderPreprocessor::ShaderPreprocessor() {}

std::stringstream ShaderPreprocessor::loadFromRelative(const std::string& path, const std::string& base) const {
    std::string fullpath = base + path;
    std::ifstream filestream;
    std::stringstream ret;
    filestream.open(fullpath);
    if (filestream.is_open()) {
        ret << filestream.rdbuf();
        return ret;
    }
    filestream.close();
    for (const std::string& ext : AcceptedExtensions) {
        filestream.open(fullpath + ext);
        if (filestream.is_open()) {
            ret << filestream.rdbuf();
            return ret;
        }
        filestream.close();
    }
    throw ShaderPreprocessorError(Codes::FILE_NOT_FOUND, base + path);
}

std::stringstream ShaderPreprocessor::loadFromDirectory(const std::string& name) const {
    std::string fullpath;
    std::ifstream filestream;
    std::stringstream ret;
    for (const std::string& path : Directories) {
        fullpath = path + name;
        filestream.open(fullpath);
        if (filestream.is_open()) {
            ret << filestream.rdbuf();
            return ret;
        }
        filestream.close();
        for (const std::string& ext : AcceptedExtensions) {
            filestream.open(fullpath + ext);
            if (filestream.is_open()) {
                ret << filestream.rdbuf();
            }
            filestream.close();
        }
    }
    throw ShaderPreprocessorError(Codes::FILE_NOT_FOUND, name);
}

std::stringstream ShaderPreprocessor::loadFromBuiltIn(const std::string& name) const {
    std::stringstream ret;
    decltype(BuiltIns)::const_iterator builtin_it = BuiltIns.find(name);
    if (builtin_it != BuiltIns.end()) {
        ret << builtin_it->second;
        return ret;
    }
    throw ShaderPreprocessorError(Codes::BUILTIN_NOT_FOUND, name);
}

std::stringstream ShaderPreprocessor::load(const std::string& path) const {
    try {
        return loadFromRelative(path);
    } catch (ShaderPreprocessorError e) {
        if (e.code == Codes::FILE_NOT_FOUND) {
            return loadFromDirectory(path);
        } else {
            throw e;
        }
    }
}

void ShaderPreprocessor::process(std::stringstream& code, ShaderUnit& shad) const {
    // std::stringstream ret;
    std::string line;
    enum Context : uint8_t {
        Code,
        BlockComment,
        FalseBranch
    };
    Context currcontext = Code;
    const size_t keySize = Keywords.size();
    std::vector<size_t> keyIdxs(keySize);
    std::vector<std::string> keys(keySize);
    std::vector<PreprocFunc> keyFuncs(keySize);
    unsigned int idx = 0;
    for (const auto& kv : Keywords) {
        keys[idx] = kv.first;
        keyFuncs[idx] = kv.second;
        idx++;
    }
    size_t linecom, blockstart, blockend;
    int lineidx = -1;
    bool loaded;
    // Use linked list to count char len of each line
    while(std::getline(code, line)) {
        lineidx++;
        keyIdxs = std::vector<size_t>(keySize, std::string::npos);
        linecom = line.find("//");
        blockstart = line.find("/*");
        blockend = line.find("*/");
        loaded = false;
        switch (currcontext) {
        case Code:
        // Figure out how to activate new contexts
        // while also handling mid-line context switching
            if (linecom != std::string::npos) {
                size_t preprocIdx;
                for (int i = 0; i < keySize; ++i) {
                    preprocIdx = line.find(keys[i]);
                    keyIdxs[i] = (preprocIdx < linecom) ? preprocIdx: std::string::npos;
                }

            } else if (blockstart != std::string::npos) {
                size_t preprocIdx;
                for (int i = 0; i < keySize; ++i) {
                    preprocIdx = line.find(keys[i]);
                    keyIdxs[i] = (preprocIdx < blockstart) ? preprocIdx: std::string::npos;
                } 
            } else {
                for (int i = 0; i < keySize; ++i) {
                    keyIdxs[i] = line.find(keys[i]);
                }
            }

            for (int i = 0; i < keySize; ++i) {
                if (keyIdxs[i] != std::string::npos) {
                    keyFuncs[i](*this, shad, 0, line);
                    loaded = true;
                    break;
                }
            }
            if (!loaded) {
                shad.preprocSource << line << "\n";
            }
            break;
        case BlockComment:
            std::cout << "In a comment!\n";
            if (blockend == std::string::npos) {
                shad.preprocSource << line << "\n";
                loaded = true;
            } else {
                size_t newpos = code.tellg() - line.size() + blockend;
                line.resize(blockend);
                shad.preprocSource << line;
                loaded = true;
                code.seekg(newpos);
                continue;
            }
            break;
        case FalseBranch:
            shad.preprocSource << line << "\n";
            loaded = true;
            break;
        default:
            shad.preprocSource << line << "\n";
            loaded = true;
            break;
        }
    }
}

ShaderUnit ShaderPreprocessor::process(std::stringstream& code) const {
    ShaderUnit ret;
    process(code, ret);
    return ret;
}

// Shader class
Shader::Shader(const std::stringstream& vertexShader, const std::stringstream& fragmentShader) {
    std::string vertexCode, fragmentCode;
    vertexCode = vertexShader.str();
    fragmentCode = fragmentShader.str();

    const char* vShaderCode = vertexCode.c_str();
    const char* fShaderCode = fragmentCode.c_str();

    unsigned int vertex, fragment; 
    int success;
    char infoLog[512];

    vertex = glCreateShader(GL_VERTEX_SHADER);
    glShaderSource(vertex, 1, &vShaderCode, NULL);
    glCompileShader(vertex);
    glGetShaderiv(vertex, GL_COMPILE_STATUS, &success);
    if(!success) {
        glGetShaderInfoLog(vertex, 512, NULL, infoLog);
        throw ShaderError("Could not compile Vertex shader", infoLog);
    }

    fragment = glCreateShader(GL_FRAGMENT_SHADER);
    glShaderSource(fragment, 1, &fShaderCode, NULL);
    glCompileShader(fragment);
    glGetShaderiv(fragment, GL_COMPILE_STATUS, &success);
    if(!success) {
        glGetShaderInfoLog(fragment, 512, NULL, infoLog);
        throw ShaderError("Could not compile Fragment shader", infoLog);
    }

    id = glCreateProgram();
    glAttachShader(id, vertex);
    glAttachShader(id, fragment);
    glLinkProgram(id);
    glGetProgramiv(id, GL_LINK_STATUS, &success);
    if (!success) {
        glGetProgramInfoLog(id, 512, NULL, infoLog);
        throw ShaderError("Could not link shader program", infoLog);
    }

    glDeleteShader(vertex);
    glDeleteShader(fragment);
}

Shader::Shader() { id = 0x0; }

Shader::Shader(const Shader& cpy_shader) { id = cpy_shader.id; }

Shader& Shader::operator=(const Shader& cpy_shader) { id = cpy_shader.id; return *this; }

void Shader::use() const{
    glUseProgram(id);
}

void Shader::setBool(const std::string& name, bool value) const {
    glUniform1i(glGetUniformLocation(id, name.c_str()), (int)value);
}

void Shader::setInt(const std::string& name, int value) const {
    glUniform1i(glGetUniformLocation(id, name.c_str()), (int)value);
}

void Shader::setFloat(const std::string& name, float value) const {
    glUniform1f(glGetUniformLocation(id, name.c_str()), (int)value);
}

void Shader::setMat4(const std::string& name, glm::mat4& value) const {
    glUniformMatrix4fv(glGetUniformLocation(id, name.c_str()), 1, GL_FALSE, glm::value_ptr(value));
}

}