#include "Shader.h"

// #define _PREPROC_FUNC_PARAM_NAMES_ (const ShaderPreprocessor& proc, ShaderUnit& shad, unsigned int linenum, const std::string& line, std::vector<uint8_t>& params)

namespace NB{

File::FilePath 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 {
            bool spec_char_found = false;
            for (const char sc : allowed_special_chars) {
                if (c == sc) {
                    path.back() += c;
                    spec_char_found = true;
                    break;
                }
            }
            if (!spec_char_found) {
                throw std::runtime_error("'" + name + "' is not valid filepath.");
            }
        }
    }
    File::FilePath ret;
    for (const auto& tk : path) {
        if (tk == path.back()) {
            size_t period = tk.find(".");
            if (period == std::string::npos) {
                ret.basename = tk;
                ret.ext = "";
            } else {
                ret.basename = tk.substr(0, period);
                ret.ext = tk.substr(period);
            }
        } else {
            ret.dir += 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::stringstream ret;
    if (file != "") {
        ret << "In file " << file;
        if (line >= 0) {
            ret << " at line " << line;
        }
        ret << ":\n\t";
    }
    ret << msg;
    if (shad != nullptr) {
        ret << " with shader error: " << shad;
    }
    return ret.str();
}

// ShaderPreprocessor class
std::string ShaderPreprocessor::TokenName(const ShaderPreprocessor::TokenType& x) {
    switch(x) {
        case TK:
            return "Token";
            break;
        case DR:
            return "Directive";
            break;
        case WS:
            return "Whitespace";
            break;
        case NL:
            return "NewLine";
            break;
        case LC:
            return "LineComment";
            break;
        case BC:
            return "BlockComment";
            break;
        default:
            return "Unrecognized";
            break;
    }
}

ShaderPreprocessor::ShaderPreprocessor() {}

File ShaderPreprocessor::loadFromBase(const std::string path, const std::string base) const {
    File ret;
    std::ifstream filestream;
    filestream.open(base + path);
    if (filestream.is_open()) {
        ret.path = get_file_path(base + path);
        ret.src << filestream.rdbuf();
        return ret;
    }
    filestream.close();
    std::string ext;
    for (const auto& kv : AcceptedExtensions) {
        ext = kv.first;
        filestream.open(base + path + ext);
        if (filestream.is_open()) {
            ret.path = get_file_path(base + path + ext);
            ret.src << filestream.rdbuf();
            return ret;
        }
        filestream.close();
    }
    throw ShaderPreprocessorError(Codes::FILE_NOT_FOUND, base + path);
}

File ShaderPreprocessor::loadFromDirectories(const std::string name) const {
    File ret;
    std::ifstream fstream;
    for (const std::string& path : Directories) {
        fstream.open(path + name);
        if (fstream.is_open()) {
            ret.path = get_file_path(path + name);
            ret.src << fstream.rdbuf();
            return ret;
        }
        fstream.close();
    }
    throw ShaderPreprocessorError(Codes::FILE_NOT_FOUND, name);
}

File ShaderPreprocessor::load(const std::string path, bool builtin_first) const {
    if(builtin_first) {
        return load_BuiltInFirst(path);
    }
    return load_FilesFirst(path);
}

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

File ShaderPreprocessor::load_FilesFirst(const std::string path) const {
    try {
        return loadFromBase(path);
    } catch (ShaderPreprocessorError e) {
        if (e.code == Codes::FILE_NOT_FOUND) {
            try {
                return loadFromDirectories(path);
            } catch (ShaderPreprocessorError f) {
                if (f.code == Codes::FILE_NOT_FOUND) {
                    try {
                        return loadFromBuiltIn(path);
                    } catch(ShaderPreprocessorError g) {
                        if (g.code == Codes::BUILTIN_NOT_FOUND) {
                            throw ShaderPreprocessorError(Codes::FILE_NOT_FOUND, path);
                        } else { throw g; }
                    }
                } else { throw f; }
            }
        } else { throw e; }
    }
}

File ShaderPreprocessor::load_BuiltInFirst(const std::string path) const {
    try {
        return loadFromBuiltIn(path);
    } catch (ShaderPreprocessorError f) {
        if (f.code == Codes::BUILTIN_NOT_FOUND) {
            try {
                return loadFromDirectories(path);
            } catch (ShaderPreprocessorError e) {
                if (e.code == Codes::FILE_NOT_FOUND) {
                    return loadFromBase(path);
                } else {
                    throw e;
                }
            }
        } else {
            throw f;
        }
    }
    
}

std::vector<ShaderPreprocessor::Token> ShaderPreprocessor::tokenize(const std::string& code) const {
    enum State {
        FSlash,
        WhiteSpace,
        LineComment,
        BlockComment,
        BlockCommentEndStar,
        Directive,
        Token
    };
    
    std::vector<ShaderPreprocessor::Token> tks;
    
    std::string token = "";
    State state = WhiteSpace;
    for(char c : code) {
        if (c==13) {
            continue;
        }
        switch(state) {
            case WhiteSpace:
                if (c=='/') {
                    if (token != "") { tks.push_back({WS, token}); }
                    token = c;
                    state = FSlash;
                } else if (c=='#') {
                    if (token != "") { tks.push_back({WS, token}); }
                    token = c;
                    state = Directive;
                } else if (c=='\n') {
                    if (token != "") { tks.push_back({WS, token}); }
                    tks.push_back({NL, "\n"});
                    token = "";
                    state = WhiteSpace;
                } else if (std::isblank(c)) {
                    token += c;
                } else {
                    if (token != "") { tks.push_back({WS, token}); }
                    token = c;
                    state = Token;
                }
                break;
            case FSlash:
                if (c=='/') {
                    token += c;
                    state = LineComment;
                } else if (c=='*') {
                    token += c;
                    state = BlockComment;
                } else if (c=='\n') {
                    tks.push_back({TK, token});
                    tks.push_back({NL, "\n"});
                    token = "";
                    state = WhiteSpace;
                } else if (std::isblank(c)) {
                    tks.push_back({TK, token});
                    token = c;
                }else {
                    token += c;
                    state = Token;
                }
                break;
            case LineComment:
                if (c=='\n') {
                    tks.push_back({LC, token});
                    tks.push_back({NL, "\n"});
                    token = "";
                    state = WhiteSpace;
                } else {
                    token += c;
                }
                break;
            case BlockComment:
                token += c;
                if (c=='*') {
                    state = BlockCommentEndStar;
                }
                break;
            case BlockCommentEndStar:
                token += c;
                if (c=='/') {
                    tks.push_back({BC, token});
                    token = "";
                    state = WhiteSpace;
                } else {
                    state = BlockComment;
                }
                break;
            case Directive:
                if (c=='\n') {
                    tks.push_back({DR, token});
                    tks.push_back({NL, "\n"});
                    token = "";
                    state = WhiteSpace;
                } else if (c=='/') {
                    tks.push_back({DR, token});
                    token = c;
                    state = FSlash;
                } else {
                    token += c;
                }
                break;
            case Token:
                if (c=='\n') {
                    tks.push_back({TK, token});
                    tks.push_back({NL, "\n"});
                    token = "";
                    state = WhiteSpace;
                } else if (std::isblank(c)) {
                    tks.push_back({TK, token});
                    token = c;
                    state = WhiteSpace;
                } else if (c=='/') {
                    tks.push_back({TK, token});
                    token = c;
                    state = FSlash;
                }else {
                    token += c;
                }
                break;
            default:
                break;
        }
        
    }

    switch(state) {
        case WhiteSpace:
            if (token != "") { tks.push_back({WS, token}); }
        case FSlash:
            tks.push_back({TK, token});
            break;
        case LineComment:
            tks.push_back({LC, token});
            break;
        case BlockComment:
        case BlockCommentEndStar:
            tks.push_back({BC, token});
            break;
        case Directive:
            tks.push_back({DR, token});
            break;
        case Token:
            tks.push_back({TK, token});
            break;
        default:
            break;
    }

    return tks;
}

ShaderUnit& ShaderPreprocessor::preprocess(const std::string& code, ShaderUnit& shad) const {
    typedef ShaderPreprocessor::Token Token;
    std::vector<Token> tks = tokenize(code);
    
    for (int i{0}; i < tks.size(); ++i) {
        switch(tks[i].first) {
            case DR:
                directive_dispatch(shad, tks[i].second);
                break;
            case TK:
            case NL:
            case LC:
            case BC:
            case WS:
            default:
                shad.preprocSource += tks[i].second;
                break;
        }
    }
    
    if (shad.vMajor == 0 && shad.vMinor == 0) {
        shad.vMajor = 1;
        shad.vMinor = 10;
    }
    return shad;
}

ShaderUnit ShaderPreprocessor::preprocess(File file, GLenum shader_type) const {
    ShaderUnit ret;
    ret.file = file;
    if (shader_type) {
        ret.type = shader_type;
    } else {
        decltype(AcceptedExtensions.begin()) find_type = AcceptedExtensions.find(file.path.ext);
        if (find_type != AcceptedExtensions.end()) {
            ret.type = find_type->second;
        }
    }
    preprocess(file.src.str(), ret);
    return ret;
}

ShaderUnit ShaderPreprocessor::preprocess(const std::string& code, GLenum shader_type) const {
    File local;
    local.path = File::FilePath{"live:", "live", ".shad"};
    return preprocess(local, shader_type);
}

ShaderProgram ShaderPreprocessor::CreateShaderProgram(std::vector<std::string> shads) const {
    std::vector<ShaderUnit> _shader_units;
    for (const auto& name : shads) {
        _shader_units.push_back(preprocess(load(name)));
    }
    return ShaderProgram::CreateShaderProgram(_shader_units);
}

ShaderProgram ShaderPreprocessor::ReloadFromFile(const ShaderProgram& rhs) const {
    std::vector<std::string> shader_names;
    shader_names.reserve(rhs._shader_units.size());
    std::string filename;
    for (const ShaderUnit& shad : rhs._shader_units) {
        filename = shad.file.path.dir + shad.file.path.basename + shad.file.path.ext;
        shader_names.emplace_back(filename);
    }
    return CreateShaderProgram(shader_names);
}

bool ShaderPreprocessor::directive_dispatch(ShaderUnit& shad, const std::string& line) const {
    StringVec dir_tks = {""};
    for (char c : line) {
        if (std::isblank(c)) {
            if (dir_tks.back() != "") {
                dir_tks.push_back("");
            }
        } else {
            dir_tks.back() += c;
        }
    }
    if (!dir_tks.size()) { return false; }

    if (dir_tks[0][0] != '#') { return false; }

    if (dir_tks[0] == "#define") {
        return preprocessor_define(shad, dir_tks, line);
    } else if (dir_tks[0] == "#version") {
        return preprocessor_version(shad, dir_tks, line);
    } else if (dir_tks[0] == "#include") {
        return preprocessor_include(shad, dir_tks, line);
    }
    return false;
}

typedef std::vector<std::string> StringVec;

bool ShaderPreprocessor::preprocessor_include(
    ShaderUnit& shad,
    const StringVec& tokens,
    const std::string& line
) const
{
    try {
        if (tokens[0] != "#include") { return false; }
    } catch (std::out_of_range e) {
        return false;
    }
    std::string path = "";
    for (const auto& tk : tokens) {
        if (tk != tokens.front()) {
            path += tk;
        }
    }
    // Add file-inclusion base + 
    // Do path cleanup + 
    // Restructure preprocessing data flow
    std::string filename = path.substr(1, path.size()-2);
    try {
        if (path[0] == '"' && path.back() == '"') {
            preprocess(load(filename).src.str(), shad);
            return true;
        } else if (path[0] == '<' && path.back() == '>') {
            preprocess(load(filename, true).src.str(), shad);
            return true;
        }
    } catch (ShaderPreprocessorError e) {
        if (e.code == Codes::FILE_NOT_FOUND) {
            std::cout << "COULD NOT FIND " << filename << ".\n";
        } else {
            throw e;
        }
    }
    return false;
}

bool ShaderPreprocessor::preprocessor_define(ShaderUnit& shad, const StringVec& tokens, const std::string& line) const {
    shad.preprocSource += line;
    try {
        if (tokens[0] != "#define") { return false; }
        try {
            shad.defines[tokens[1]] = tokens[2];
        } catch (std::length_error& f) {
            shad.defines[tokens[1]] = "";
        }
        return true;
    } catch (std::out_of_range& e) {
        return false;
    }
}

bool ShaderPreprocessor::preprocessor_version(ShaderUnit& shad, const StringVec& tokens, const std::string& line) const {
    shad.preprocSource += line;
    std::cout << "Shader version: ";
    try {
        if (tokens[0] != "#version") { return false; }
        if (tokens[1].size() == 3) {
            short vMajor = tokens[1][0]-'0';
            short vMinorA = tokens[1][1]-'0';
            short vMinorB = tokens[1][2]-'0';
            if (vMajor > 10 || vMinorA > 10 || vMinorB > 10) {
                return false;
            }
            shad.vMajor = vMajor;
            shad.vMinor = vMinorA*10 + vMinorB;
            shad.profile = AcceptedProfiles.at(tokens[2]);
            std::cout << shad.vMajor << "." << shad.vMinor << "\n";
            return true;
        } else {
            return false;
        }
    } catch (std::out_of_range& e) {
        return false;
    }
}


/* bool ShaderPreprocessor::preprocessor_uniform(
    ShaderUnit& shad,
    const std::string& type,
    const std::string& name
) const {
    std::string type_str = "";
    unsigned int type_str_len = 0;
    for (const char& c : type) {
        if (type_str_len) {
            if (std::isdigit(c)) {
                type_str += c;
                continue;
            } else if (c=='[') {
                type_str += '_';
            }else if (c == ']' || std::isblank(c)){
                continue;
            } else {
                return false;
            }
        } else {
            if (std::isalnum(c)) {
                type_str += c;
                continue;
            } else if (c == '[') {
                type_str_len = type_str.length();
                type_str += '_';
                continue;
            } else {
                return false;
            }
        }
    }
    std::string name_str = "";
    unsigned int name_str_len = 0;
    for (const char& c : name) {
        if (c == ';') {
            break;
        }
        if (name_str_len) {
            if (std::isdigit(c)) {
                type_str.insert(type_str_len, 1, c);
                type_str_len++;
                continue;
            } else if (c=='[') {
                type_str.insert(type_str_len, "_");
                type_str_len++;
                continue;
            }else if (c == ']' || std::isblank(c)){
                continue;
            } else {
                return false;
            }
        } else {
            if (std::isalnum(c)) {
                name_str += c;
                continue;
            } else if ( c == '[') {
                name_str_len = name_str.length();
                type_str.insert(type_str_len, "_");
                type_str_len++;
                continue;
            } else {
                return false;
            }
        }
    }
    shad.uniforms.push_back(UniformHandle{
        name_str,
        type_str,
        0x0,
        0x0
    });
    return true;
} */


// ShaderProgram

ShaderProgram::ShaderProgram(ShaderProgram&& rhs) {
    _shader_units = rhs._shader_units;
    _id = rhs._id;
    rhs._id = 0;
}

ShaderProgram::~ShaderProgram() {
    glDeleteProgram(_id);
}

ShaderProgram& ShaderProgram::operator=(ShaderProgram&& rhs) {
    _shader_units = rhs._shader_units;
    _id = rhs._id;
    rhs._id = 0;
    return *this;
}


ShaderProgram ShaderProgram::CreateShaderProgram(std::vector<ShaderUnit>& shaders) {
    int success;
    ShaderProgram ret(shaders);
    char infoLog[512];
    unsigned int shad_id;
    std::vector<unsigned int> shad_ids;
    shad_ids.reserve(shaders.size());
    for (auto& shad : ret._shader_units) {
        const char* source = shad.preprocSource.data();
        shad_id = glCreateShader(shad.type);
        shad_ids.emplace_back(shad_id);
        glShaderSource(shad_id, 1, &source, NULL);
        glCompileShader(shad_id);
        glGetShaderiv(shad_id, GL_COMPILE_STATUS, &success);
        if (!success) {
            glGetShaderInfoLog(shad_id, 512, NULL, infoLog);
            File::FilePath& fp = shad.file.path;
            std::string filename = fp.dir + fp.basename + fp.ext;
            // std::cout << "Could not compile '" + filename + "': " << infoLog << "\n";
            // return *ret;
            throw ShaderError("Could not compile '" + filename + "'.", infoLog);
        }
    }
    
    ret._id = glCreateProgram();
    for(auto& id : shad_ids) {
        glAttachShader(ret._id, id);
    }
    glLinkProgram(ret._id);
    glGetProgramiv(ret._id, GL_LINK_STATUS, &success);
    if (!success) {
        glGetProgramInfoLog(ret._id, 512, NULL, infoLog);
        throw ShaderError("Could not link shader program.", infoLog);
    }
    for (auto& id : shad_ids) {
        glDeleteShader(id);
    }
    return ret;
}

unsigned int ShaderProgram::id() const { return _id; }

void ShaderProgram::use() const {
    glUseProgram(_id);
}

std::vector<ShaderUnit> ShaderProgram::getShaders() const {
    return _shader_units;
}

ShaderUnit ShaderProgram::getShaders(unsigned int idx) const {
    return _shader_units[idx];
}

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

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

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

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

/* // Shader class
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);
}

*/

}