// Copyright (c) 2023 Dominic Masters
// 
// This software is released under the MIT License.
// https://opensource.org/licenses/MIT

#pragma once
#include "display/shader/ShaderStage.hpp"
#include "display/shader/IShader.hpp"
#include "assert/assert.hpp"
#include "assert/assertgl.hpp"
#include "display/Color.hpp"
#include "display/Texture.hpp"

namespace Dawn {
  typedef GLuint shadertexturebinding_t;

  enum class ShaderOpenGLVariant {
    GLSL_330_CORE
  };

  struct ShaderOpenGLParameter {
    std::string name;
    size_t offset;
    enum ShaderParameterType type;
    int32_t count;

    GLint location = -1;

    ShaderOpenGLParameter(
      const std::string &name,
      const void *offset,
      const enum ShaderParameterType type
    ) {
      this->name = name;
      this->offset = (size_t)offset;
      this->type = type;
      this->count = 1;
    }

    ShaderOpenGLParameter(
      const std::string &name,
      const void *offset,
      const enum ShaderParameterType type,
      const int32_t count
    ) {
      this->name = name;
      this->offset = (size_t)offset;
      this->type = type;
      this->count = count;
    }
  };

  template<typename T>
  class Shader : public IShader<T> {
    private:
      std::vector<std::shared_ptr<ShaderStage>> stages;
      std::vector<struct ShaderOpenGLParameter> parameters;
      enum ShaderOpenGLVariant variant;

      GLuint shaderProgram = -1;

    protected:
      virtual void getStages(
        const enum ShaderOpenGLVariant variant,
        const T *rel,
        std::vector<std::shared_ptr<ShaderStage>> &stages,
        std::vector<struct ShaderOpenGLParameter> &parameters
      ) = 0;

    public:
      /**
       * Initializes the shader, this needs to be called before the shader can
       * be used.
       */
      void init() override {
        // Determine which kind of OpenGL shader to use.
        variant = ShaderOpenGLVariant::GLSL_330_CORE;

        // Now get the stages
        T dummy;
        this->getStages(
          variant,
          &dummy,
          stages,
          parameters
        );

        // Create the shader program
        shaderProgram = glCreateProgram();
        assertNoGLError();

        // Attach all the stages
        for(auto stage : stages) {
          glAttachShader(shaderProgram, stage->id);
          assertNoGLError();
        }

        // Link and verify the program
        glLinkProgram(shaderProgram);
        assertNoGLError();

        GLint status;
        glGetProgramiv(shaderProgram, GL_LINK_STATUS, &status);
        assertNoGLError();
        assertTrue(status == GL_TRUE, "Failed to link shader program.");

        // Map parameters correctly.
        std::for_each(
          parameters.begin(),
          parameters.end(),
          [&](struct ShaderOpenGLParameter &param) {
            // Correct offset
            param.offset = param.offset - (size_t)(&dummy);
            param.location = glGetUniformLocation(
              shaderProgram,
              param.name.c_str()
            );
            assertNoGLError();
            assertTrue(
              param.location != -1,
              "Failed to get location for parameter %s.",
              param.name.c_str()
            );
          }
        );

        this->bind();
      }

      /**
       * Binds the shader as the current one, does not upload any data, somewhat
       * relies on something else uploading the data.
       */
      void bind() override {
        glUseProgram(shaderProgram);
        assertNoGLError();
      }

      /**
       * Uploads the data to the GPU.
       */
      void upload() override {
        switch(this->variant) {
          case ShaderOpenGLVariant::GLSL_330_CORE:
            for(auto param : parameters) {
              void *value = (void*)(
                ((size_t)&this->data) + param.offset
              );

              switch(param.type) {
                case ShaderParameterType::MAT4: {
                  glm::mat4 *matrix = (glm::mat4 *)value;
                  if(param.count != 1) {
                    assertUnreachable("I haven't implemented multiple mat4s");
                  }
                  glUniformMatrix4fv(
                    param.location, 1, GL_FALSE, glm::value_ptr(*matrix)
                  );
                  break;
                }
                
                case ShaderParameterType::COLOR: {
                  auto color = (Color *)value;
                  glUniform4fv(
                    param.location,
                    param.count,
                    (GLfloat*)value
                  );
                  break;
                }

                case ShaderParameterType::BOOLEAN: {
                  glUniform1iv(param.location, param.count, (GLint*)value);
                  break;
                }

                case ShaderParameterType::TEXTURE: {
                  glUniform1iv(param.location, param.count, (GLint*)value);
                  break;
                }

                default: {
                  assertUnreachable("Unsupported ShaderParameterType");
                }
              }

              assertNoGLError();
            }
            break;

          default:
            assertUnreachable("Unsupported ShaderOpenGLVariant");
        }
      }

      ~Shader() {
      }
  };
}