#ifndef __OPENGLPLAYGROUND_SHADERS_HH__
#define __OPENGLPLAYGROUND_SHADERS_HH__

#include <string>
#include <unordered_map>

#define GL_GLEXT_PROTOTYPES
#include <GL/gl.h>

#include "common.hh"

class ShaderException : public GLException {
public:
  ShaderException(std::string const& msg) : GLException(glGetError()), _msg(msg) {
  }

  virtual ~ShaderException() {}

  std::string const& getMsg() const {return _msg;}
  
  virtual std::string toString() const {
    return "ShaderException: " + _msg;
  }
  
private:
  std::string _msg;
};

class Program;

class Shader {
public:
  Shader() : _shaderID(0) {
  }
  
  Shader(std::string const& program, GLenum type) {
    _shaderID = glCreateShader(type);
    const char* const arr[] = {program.c_str()};
    glShaderSource(_shaderID, 1, arr, NULL);
    glCompileShader(_shaderID);
    
    int state;
    glGetShaderiv(_shaderID, GL_COMPILE_STATUS, &state);
    
    if (!state) {
      int logLength;
      glGetShaderiv(_shaderID, GL_INFO_LOG_LENGTH, &logLength);
      
      char *log = new char[logLength];
      glGetShaderInfoLog(_shaderID, logLength, NULL, log);
      std::string msg(log);
      delete[] log;
      
      glDeleteShader(_shaderID);
      
      throw ShaderException(msg);
    }
  }

  virtual ~Shader() {
    glDeleteShader(_shaderID);
  }

protected:
  unsigned getID() const {
    return _shaderID;
  }

  friend class Program;
 
  unsigned int _shaderID;
};
  

class VertexShader : public Shader {
public:
  VertexShader(std::string const& program) : Shader(program, GL_VERTEX_SHADER) {
  }

  virtual ~VertexShader() {
  }
};

class FragmentShader : public Shader {
public:
  FragmentShader(std::string const& program) : Shader(program, GL_FRAGMENT_SHADER) {
  }

  virtual ~FragmentShader() {
  }
};

class GeometryShader : public Shader {
public:
  GeometryShader(std::string const& program) : Shader(program, GL_GEOMETRY_SHADER) {
  }

  virtual ~GeometryShader() {
  }
};

class Program {
public:
  Program(VertexShader& vertex, FragmentShader& frag) : _geom(nullptr), _vertex(vertex), _frag(frag) {
    _progID = glCreateProgram();
    
    glAttachShader(_progID, _vertex.getID());
    glAttachShader(_progID, _frag.getID());

    glLinkProgram(_progID);
    
    int state;
    glGetProgramiv(_progID, GL_LINK_STATUS, &state);
    
    if (!state) {
      int logLength;
      glGetProgramiv(_progID, GL_INFO_LOG_LENGTH, &logLength);
      
      char *log = new char[logLength];
      glGetProgramInfoLog(_progID, logLength, NULL, log);
      std::string msg(log);
      delete[] log;
      
      glDeleteProgram(_progID);
      
      throw ShaderException(msg);
    }

    glDetachShader(_progID, _vertex.getID());
    glDetachShader(_progID, _frag.getID());
  }

  Program(GeometryShader& geom, VertexShader& vertex, FragmentShader& frag) : _geom(&geom), _vertex(vertex), _frag(frag) {
    _progID = glCreateProgram();

    glAttachShader(_progID, _geom->getID());
    glAttachShader(_progID, _vertex.getID());
    glAttachShader(_progID, _frag.getID());

    glLinkProgram(_progID);
    
    int state;
    glGetProgramiv(_progID, GL_LINK_STATUS, &state);
    
    if (!state) {
      int logLength;
      glGetProgramiv(_progID, GL_INFO_LOG_LENGTH, &logLength);
      
      char *log = new char[logLength];
      glGetProgramInfoLog(_progID, logLength, NULL, log);
      std::string msg(log);
      delete[] log;
      
      glDeleteProgram(_progID);
      
      throw ShaderException(msg);
    }

    glDetachShader(_progID, _geom->getID());
    glDetachShader(_progID, _vertex.getID());
    glDetachShader(_progID, _frag.getID());
  }
  
  ~Program() {
    glDeleteProgram(_progID);
  }

  void use() const {
    glUseProgram(_progID);
  }

  int getUniformLocation(std::string const& name) const {
    auto search = _uniformLocCache.find(name);
    if (search != _uniformLocCache.end())
      return search->second;
   
    int ret = glGetUniformLocation(_progID, name.c_str());
    if (ret == -1)
      throw GLException(glGetError());
    _uniformLocCache.emplace(name, ret);
    return ret;
  }

  int getAttribLocation(std::string const& name) const {
    auto search = _attribLocCache.find(name);
    if (search != _attribLocCache.end())
      return search->second;
    
    int ret = glGetAttribLocation(_progID, name.c_str());
    if (ret == -1)
      throw GLException(glGetError());
    _attribLocCache.emplace(name, ret);
    return ret;
  }
  
private:
  GeometryShader* _geom;
  VertexShader& _vertex;
  FragmentShader& _frag;
  unsigned _progID;

  mutable std::unordered_map<std::string, GLint> _uniformLocCache, _attribLocCache;
};

#endif