#include <GL/glew.h>
#include <tinyxml2.h>
#include <cstdio>
#include <memory.h>
#include "shader.h"
#include "../render.h"
#include <log.h>

using namespace tinyxml2;

const std::string projectionMatrixName = "projectionMatrix";
const std::string viewMatrixName = "viewMatrix";
const std::string modelMatrixName = "modelMatrix";
const std::string normalMatrixName = "normalMatrix";

namespace render{

shader::shader()
{
    _attachedShaders = 0;
    program = 0;
    _linked = 0;
}

shader::~shader()
{
    free();
}

int shader::getUniformLocation(const std::string &name) const
{
    if(!_linked)
        throw std::runtime_error("shader not linked");
    bind();
    auto it = _uniforms.lower_bound(name);

    if(it != _uniforms.end() && !(_uniforms.key_comp()(name, it->first)))
        return it->second->_id;

    throw std::runtime_error(name+" uniform not found");
}

uniform &shader::operator[](const std::string &name) const
{
    bind();
    auto it = _uniforms.lower_bound(name);
    if(it != _uniforms.end() && !(_uniforms.key_comp()(name, it->first)))
        return *(it->second);
    throw std::runtime_error(name+" uniform not found");
}

void shader::_alloc()
{
    program = glCreateProgram();
    _addToList();
}

bool shader::load(const std::string &path)
{
    XMLDocument doc;

    _log::out("Loading shader %s\n", path.c_str());

    if(doc.LoadFile(path.c_str()) != XML_SUCCESS)
        return 0;

    std::string p = path.substr(0, path.rfind('/')+1);

    const char *binName = 0;
    uint types[3] = {SHADER_VERTEX, SHADER_FRAGMENT, SHADER_GEOMETRY};
    const char *names[] = {"vertex", "fragment", "geometry"};
    XMLElement *root = doc.RootElement();

    if(binName = root->Attribute("binary"))
        if(loadBinary(p+binName))
            return 1;

    for(int i = 0; i < 3; i++)
        if(XMLElement *elem = root->FirstChildElement(names[i]))
        {
            if(elem->Attribute("path"))
                attach(p+elem->Attribute("path"), types[i]);
            else
                attach(p+elem->GetText(), types[i]);
        }

    linkProgram(binName != nullptr);
    validate();

    _getAllUniforms();

    if(binName)
        saveBinary(p+binName);

    return 1;
}

bool shader::attach(const std::string &path, uint s)
{
    if(_linked) return 0;
    if(s > SHADER_GEOMETRY) return 0;
    int length;
    const char *names[] = {"vertex", "fragment", "geometry"};
    _log::out("Loading %s shader: %s ", names[s], path.c_str());
    FILE *src = fopen(path.c_str(), "rt");
    try
    {
        if(!src)
            throw "Failed to open file";
        fseek(src, 0, 2);
        if(!(length = ftell(src)))
            throw "File is empty";
        fseek(src, 0, 0);
        std::string content;
        content.resize(length+1);
        fread(&content[0], 1, length, src);
        fclose(src);
        attachSource(content.data(), s);
    }
    catch(const char* msg)
    {
        _log::out("[FAIL]\n%s\n", msg);
        if(src)
            fclose(src);
        return 0;
    }
    _log::out("[ OK ]\n");
    return 1;
}

bool shader::attachSource(const char *src, uint s)
{
    if(_linked) return 0;
    if(s > SHADER_GEOMETRY) return 0;
    uint types[3] = {GL_VERTEX_SHADER, GL_FRAGMENT_SHADER, GL_GEOMETRY_SHADER};
    const char *names[] = {"vertex", "fragment", "geometry"};
    if(!program)
        _alloc();

    if(!glIsShader(shaders[s]))
        shaders[s] = glCreateShader(types[s]);
    else
        if(_attachedShaders&(1<<s))
        {
            glDetachShader(program, shaders[s]);
            _attachedShaders &= ~(1<<s);
        }
    glShaderSource(shaders[s], 1, &src, 0);
    glCompileShader(shaders[s]);

    int status = 0;
    glGetShaderiv(shaders[s], GL_COMPILE_STATUS, &status);
    if(!status)
    {
        _printSError(shaders[s]);
        return 0;
    }

    glAttachShader(program, shaders[s]);
    _attachedShaders |= (1<<s);
    return 1;
}

bool shader::linkProgram(bool binaryHint)
{
    if(_linked)
        return 0;
    if(binaryHint)
        glProgramParameteri(program, GL_PROGRAM_BINARY_RETRIEVABLE_HINT, GL_TRUE);
    glLinkProgram(program);
    int lenth;

    _log::out("Linking shader program ");
    glGetProgramiv(program, GL_LINK_STATUS, &lenth);
    if(!lenth)
    {
        _log::out("[FAIL]\n");
        _printPError(program);
        return 0;
    }
    _log::out("[ OK ]\n");
    _linked = 1;
    return 1;
}

bool shader::validate()
{
    if(!_linked) return 0;
    glValidateProgram(program);
    int s = -1;
    glGetProgramiv(program, GL_VALIDATE_STATUS, &s);
    _log::out("Validate shader %s\n", (s==1)?"[ OK ]":"[FAIL]");
    return s;
}

bool shader::loadBinary(const std::string &path)
{
    if(_linked)
        throw std::runtime_error("shader already linked");
    int formatsCount = 0, length = 0;
    GLenum format;
    _log::out("loading binary shader: %s ", path.c_str());
    FILE *f;
    try
    {
        if(!(f = fopen(path.c_str(), "rb")))
            throw std::runtime_error("Failed to open file");
        fseek(f, 0, 2);
        if(!(length = ftell(f)))
            throw std::runtime_error("File is empty");
        fseek(f, 0, 0);
        length -= sizeof(GLenum);
        char *data = new char[length];
        fread(&format, sizeof(GLenum), 1, f);
        fread(data, length, 1, f);

        _alloc();
        glProgramBinary(program, format, data, length);
        delete [] data;
        glGetProgramiv(program, GL_LINK_STATUS, &length);
        if(!length)
            throw std::runtime_error("Some error");
        _linked = 1;
        fclose(f);
    }
    catch(const std::runtime_error &ex)
    {
        _log::out("[FAIL]\n%s\n", ex.what());
        _printPError(program);
        if(f) fclose(f);
        return 0;
    }
    _getAllUniforms();
    _log::out("[ OK ]\n");
    return 1;
}

bool shader::saveBinary(const std::string &path)
{
    if(!_linked)
        return 0;
    int length = 0;
    glGetProgramiv(program, GL_PROGRAM_BINARY_LENGTH, &length);
    if(!length)
        return 0;
    char *data = new char[length];
    GLenum format;
    glGetProgramBinary(program, length, 0, &format, data);
    FILE *out = fopen(path.c_str(), "wb");
    if(out)
    {
        fwrite(&format, sizeof(GLenum), 1, out);
        fwrite(data, length, 1, out);
        fclose(out);
    }
    delete [] data;
    return 1;
}

void shader::_getAllUniforms()
{
    int count, size, length;
    GLenum type;
    char name[256];
    glGetProgramiv(program, GL_ACTIVE_UNIFORMS, &count);
    for(int i = 0; i < count; i++)
    {
        glGetActiveUniform(program, i, 256, &length, &size, &type, name);
        _uniforms[std::string(name, length)] = new uniform(i, this);
    }
}

void shader::bind() const
{
    if(!_linked) return;
    render::bindShader(this);
}

void shader::unbind() const
{
    render::bindShader(0);
}

void shader::_printSError(GLuint i)
{
    int log = 0;
    glGetShaderiv(i, GL_INFO_LOG_LENGTH, &log);
    if(!log)
        return;
    char *mess = new char[log];
    glGetShaderInfoLog(i, log, 0, mess);
    _log::out("\t%s\n", mess);
    delete [] mess;
}

void shader::_printPError(GLuint i)
{
    int log = 0;
    glGetProgramiv(i, GL_INFO_LOG_LENGTH, &log);
    if(!log)
        return;
    char *mess = new char[log];
    glGetProgramInfoLog(i, log, 0, mess);
    _log::out("\t%s\n", mess);
    delete [] mess;
}

void shader::free()
{
    unbind();
    if(glIsProgram(program))
    {
        for(int i = 0 ; i < 3; i++)
            if(_attachedShaders&(1<<i))
                glDetachShader(program, shaders[i]);
        glDeleteProgram(program);
        program = 0;
        _attachedShaders = 0;
    }
    for(int i = 0; i < 3; i++)
        if(glIsShader(shaders[i]))
        {
            glDeleteShader(shaders[i]);
            shaders[i] = 0;
        }
    _removeFromList();

    for(const auto &i: _uniforms)
        delete i.second;
    _uniforms.clear();
}

uint shader::getProgram() const
{
    return program;
}

void shader::loadProjectionMatrix(const float *mat) const
{
    uniform4x4(projectionMatrixName, 1, mat);
}

void shader::loadModelMatrix(const float *mat) const
{
    uniform4x4(modelMatrixName, 1, mat);
}

void shader::loadViewMatrix(const float *mat) const
{
    uniform4x4(viewMatrixName, 1, mat);
}

void shader::loadNormalMatrix(const float *mat) const
{
    uniform3x3(normalMatrixName, 1, mat);
}

void shader::bindAttribLocation(const std::string &name, uint val)
{
    glBindAttribLocation(program, val, name.c_str());
}

uint shader::getAttribLocation(const std::string &name)
{
    if(_linked)
        return glGetAttribLocation(program, name.c_str());
    return -1;
}

void shader::uniform1f(const std::string &name, float v1)
{
    flush();
    glUniform1f(getUniformLocation(name), v1);
}

void shader::uniform2f(const std::string &name, float v1, float v2)
{
    flush();
    glUniform2f(getUniformLocation(name), v1, v2);
}

void shader::uniform3f(const std::string &name, float v1, float v2, float v3)
{
    flush();
    glUniform3f(getUniformLocation(name), v1, v2, v3);
}

void shader::uniform4f(const std::string &name, float v1, float v2, float v3, float v4)
{
    flush();
    glUniform4f(getUniformLocation(name), v1, v2, v3, v4);
}

///=================================================

void shader::uniform1i(const std::string &name, int v1)
{
    flush();
    glUniform1i(getUniformLocation(name), v1);
}

void shader::uniform2i(const std::string &name, int v1, int v2)
{
    flush();
    glUniform2i(getUniformLocation(name), v1, v2);
}

void shader::uniform3i(const std::string &name, int v1, int v2, int v3)
{
    flush();
    glUniform3i(getUniformLocation(name), v1, v2, v3);
}

void shader::uniform4i(const std::string &name, int v1, int v2, int v3, int v4)
{
    flush();
    glUniform4i(getUniformLocation(name), v1, v2, v3, v4);
}

void shader::uniform4x4(const std::string &name, uint count, const float *v) const
{
    glUniformMatrix4fv(getUniformLocation(name), count, 0, v);
}

void shader::uniform3x3(const std::string &name, uint count, const float *v) const
{
    glUniformMatrix3fv(getUniformLocation(name), count, 0, v);
}

///=============================================================================

uniform::uniform(uint id, shader *root)
{
    _id = id;
    _root = root;
}

void uniform::operator=(float v) const noexcept
{
    _root->bind();
    flush();
    glUniform1f(_id, v);
}

void uniform::operator=(const vec2f &v) const noexcept
{
    _root->bind();
    flush();
    glUniform2f(_id, v.x, v.y);
}

void uniform::operator=(const vec3f &v) const noexcept
{
    _root->bind();
    flush();
    glUniform3f(_id, v.x, v.y, v.z);
}

void uniform::operator=(const vec4f &v) const noexcept
{
    _root->bind();
    flush();
    glUniform4f(_id, v.x, v.y, v.z, v.y);
}

void uniform::operator=(const glm::mat4 &v) const noexcept
{
    _root->bind();
    flush();
    glUniformMatrix4fv(_id, 1, 0, &v[0][0]);
}

}
