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

// 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) const {
/*     std::stringstream ret;
    std::string line;
    enum State : unsigned char {
        Start,
        Token,
        PreProcToken,
        Number,
        WhiteSpace,
        FSlash,
        Operator,
        BlockComment,
        LineComment
    };
    std::pair<std::string, State> const NewState = {"", Start};
    std::vector<std::pair<std::string, State>> tokens({NewState});
    State currState;
    char n;
    for (char c; code.good() && !code.eof(); code.get(c)) {
        currState = tokens.back().second;
        if (std::isblank(c) && currState != BlockComment && currState != LineComment) {
            if (c == '\n') {
                
            }
            switch (currState) {
                case WhiteSpace:
                    tokens.back().first += c;
                    continue;
                break;
                case Start:
                    tokens.back() = {"" + c, WhiteSpace};
                    continue;
                break;
                default:
                    tokens.push_back({"" + c, WhiteSpace});
                    continue;
                break;
            }
        }
        switch (currState) {
            case Start:
                if (c == '_' || std::isalpha(c)) {
                    tokens.back() = {"" + c, Token};
                    continue;
                } else if (c == '#') {
                    tokens.back() = {"#", PreProcToken};
                    continue;
                } else if (std::isdigit(c)) {
                    tokens.back() = {"" + c, Number};
                    continue;
                } else if (c == '.') {
                    code.get(n);
                    if (std::isdigit(n)) {
                        tokens.back() = {"" + c + n, Number};
                        continue;
                    } else {
                        code.unget();
                        tokens.back() = {".", Operator};
                        tokens.push_back(NewState);
                        continue;
                    }
                } else {
                    tokens.back() = {"" + c, Operator};
                    tokens.push_back(NewState);
                }
            break;
            case FSlash:
                if (c == '*') {
                    tokens.back() = {"/*", BlockComment};
                } else if (c == '/') {
                    tokens.back() = {"//", LineComment};
                } else {
                    tokens.back().second = Operator;
                }
                continue;
            break;
            case BlockComment:
                tokens.back().first += c;
                code.get(n);
                tokens.back().first += n;
                    if (c == '*' && n == '/') {
                    tokens.push_back(NewState);
                }
            break;
            case LineComment:
                tokens.back().first += c;
                if (c == '\n') {
                    tokens.push_back(NewState);
                }
            break;
        }
    }

    return ret; */
    std::stringstream ret;
    const std::vector<std::string> PreprocessorKeywords = {"#include"};

    enum State : unsigned char {
        Symbol,
        Directive,
        BlockCommentStart,
        BlockCommentEnd,
        LinecommentStart,
        Whitespace,
        Pound,
        SingleQuote,
        DoubleQuote,
        OpenAngle,
        CloseAngle,
        FSlash,
        Semicolon,
        Star,
        Newline,
        Operator
    };
    std::vector<std::pair<State, std::string>> stack = {};
    for (char c; code.good() && !code.eof(); code.get(c)) {
        State currstate = stack.back().first;
        bool editstack = false;
        switch (c) {
            case '#':
                stack.push_back({Pound, "#"});
                break;
            case '"':
                stack.push_back({DoubleQuote, "\""});
                break;
            case '\'':
                stack.push_back({SingleQuote, "'"});
                break;
            case '<':
                stack.push_back({OpenAngle, "<"});
                break;
            case '>':
                stack.push_back({CloseAngle, ">"});
                break;
            case '/':
                stack.push_back({FSlash, "/"});
                break;
            case '*':
                stack.push_back({Star, "*"});
                break;
            case '\n':
                stack.push_back({Newline, "\n"});
                break;
            default:
                if (isblank(c)) {
                    stack.push_back({Whitespace, std::string(1, c)});
                } else if (isalnum(c)) {
                    if (currstate == Symbol) {
                        stack.back().second += c;
                    } else {
                        stack.push_back({Symbol, std::string(1, c)});
                    }
                } else {
                    stack.push_back({Operator, std::string(1, c)});
                }
                break;
        }

        auto endIt = stack.end();
        State secondlast = endIt[-2].first;
        switch (endIt[-1].first) {
            case FSlash:
                if (secondlast == FSlash) {
                    stack.erase(endIt-2, endIt);
                    stack.push_back({LinecommentStart, "//"});
                } else if (secondlast == Star) {
                    stack.erase(endIt-2, endIt);
                    stack.push_back({BlockCommentEnd, "*/"});
                }
                break;
            case Star:
                if (secondlast == FSlash) {
                    stack.erase(endIt-2, endIt);
                    stack.push_back({BlockCommentStart, "/*"});
                }
                break;
            case Newline:
                for (auto i : stack) {

                }

        }
    }
}

std::stringstream ShaderPreprocessor::process(std::stringstream& code) const {
    std::stringstream ret;
    std::string line;
    const std::vector<std::string> preprocKeys = {
        "#include",
        "#define",
        "#ifdef",
        "#ifndef",
        "#endif"
    };
    enum Context : uint8_t {
        Code,
        BlockComment,
        FalseBranch
    };
    Context currcontext = Code;
    const size_t keysListSize = preprocKeys.size();
    std::vector<size_t> keysIdx(keysListSize);
    size_t linecom, blockstart, blockend;
    // Use linked list to count char len of each line
    while(std::getline(code, line)) {
        switch (currcontext) {
        case Code:
            if (linecom != std::string::npos) {
                size_t preprocIdx;
                for (int i = 0; i < keysListSize; ++i) {
                    preprocIdx = line.find(preprocKeys[i]);
                    keysIdx[i] = (preprocIdx < linecom) ? preprocIdx: std::string::npos;
                }
            } else if (blockstart != std::string::npos) {
                size_t preprocIdx;
                for (int i = 0; i < keysListSize; ++i) {
                    preprocIdx = line.find(preprocKeys[i]);
                    keysIdx[i] = (preprocIdx < blockstart) ? preprocIdx: std::string::npos;
                } 
            } else {
                for (int i = 0; i < keysListSize; ++i) {
                    keysIdx[i] = line.find(preprocKeys[i]);
                }
            }

            if (keysIdx)
            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;
        }
    }
}

std::stringstream ShaderPreprocessor::process(std::string name) const {
    std::stringstream code = load(name);
    return process(code);
}

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

}