#include "Shader.h"

namespace NB{

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

static void preprocessor_include(const ShaderPreprocessor& proc, ShaderUnit& shad, std::stringstream& stream, unsigned int linenum, const std::string& line) {
    std::stringstream lstream(line);
    std::vector<std::string> tokens;
    std::string token;
    while (std::getline(lstream, token, ' ')) {
        tokens.push_back(token);
    }
    if (tokens[0] != "#include") { return; }
    shad.preprocSource << proc.include(" ", "", false);
}

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

std::stringstream ShaderPreprocessor::include(const std::string& name, const std::string& base, bool relativeFirst) const {
    if (relativeFirst) {
        try {
            return loadFromRelative(name, base);
        } catch (ShaderPreprocessorError e) {
            if (e.code == Codes::FILE_NOT_FOUND) {
                try {
                    return loadFromDirectory(name);
                } catch (ShaderPreprocessorError f) {
                    if (f.code == Codes::FILE_NOT_FOUND) {
                        return loadFromBuiltIn(name);
                    } else {
                        throw f;
                    }
                }
            } else {
                throw e;
            }
        }
    }
    try {
        return loadFromDirectory(name);
    } catch (ShaderPreprocessorError e) {
        if (e.code == Codes::FILE_NOT_FOUND) {
            try {
                return loadFromBuiltIn(name);
            } catch (ShaderPreprocessorError f) {
                if (f.code == Codes::BUILTIN_NOT_FOUND) {
                    return loadFromRelative(name, base);
                } else {
                    throw f;
                }
            }
        } else {
            throw e;
        }
    }
}

std::stringstream 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;
    // Use linked list to count char len of each line
    while(std::getline(code, line)) {
        keyIdxs = std::vector<size_t>(keySize, std::string::npos);
        linecom = 0;
        blockstart = 0;
        blockend = 0;
        switch (currcontext) {
        case Code:
            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, ret, );
                }
            } 

            break;
        case BlockComment:
            if (blockend == std::string::npos) {
                ret << line;
            } else {
                size_t newpos = code.tellg() - line.size() + blockend;
                line.resize(blockend);
                ret << line;
                code.seekg(newpos);
                continue;
            }
            break;
        case FalseBranch:
            ret << line;
            break;
        default:
            ret << line;
            break;
        }
    }
}

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

}