#include <pch.hpp>

#include <core.hpp>
#include <opengl/opengl_shader.hpp>
#include <instrumentor.hpp>
#include <core/time.hpp>
#include <logging.hpp>

#include <shaderc/shaderc.hpp>
#include <spirv_cross/spirv_cross.hpp>
#include <spirv_cross/spirv_glsl.hpp>
#include <glm/gtc/type_ptr.hpp>
#include <glm/glm.hpp>
#include <filesystem>
#include <glad/glad.h>
//#include <alloca.h>

namespace OpenEngine {
    
    namespace Utils {

        static GLenum ShaderTypeFromString(const std::string& type)
        {
            if (type == "vertex")
                return GL_VERTEX_SHADER;
            if (type == "fragment" || type == "pixel")
                return GL_FRAGMENT_SHADER;

            OE_CORE_ASSERT(false, "Unknown shader type!");
            return 0;
        }

        static shaderc_shader_kind GLShaderStageToShaderC(GLenum stage)
        {
            switch (stage)
            {
                case GL_VERTEX_SHADER:   return shaderc_glsl_vertex_shader;
                case GL_FRAGMENT_SHADER: return shaderc_glsl_fragment_shader;
            }
            OE_CORE_ASSERT(false, "No");
            return (shaderc_shader_kind)0;
        }

        static const char* GLShaderStageToString(GLenum stage)
        {
            switch (stage)
            {
                case GL_VERTEX_SHADER:   return "GL_VERTEX_SHADER";
                case GL_FRAGMENT_SHADER: return "GL_FRAGMENT_SHADER";
            }
            OE_CORE_ASSERT(false, "No");
            return nullptr;
        }

        static const char* GetCacheDirectory()
        {
            // TODO: make sure the assets directory is valid
            return "assets/cache/shader/opengl";
        }

        static void CreateCacheDirectoryIfNeeded()
        {
            std::string cache_directory = GetCacheDirectory();
            if (!std::filesystem::exists(cache_directory))
                std::filesystem::create_directories(cache_directory);
        }

        static const char* GLShaderStageCachedOpenGLFileExtension(uint32_t stage)
        {
            switch (stage)
            {
                case GL_VERTEX_SHADER:    return ".cached_opengl.vert";
                case GL_FRAGMENT_SHADER:  return ".cached_opengl.frag";
            }
            OE_CORE_ASSERT(false, "No");
            return "";
        }

        static const char* GLShaderStageCachedVulkanFileExtension(uint32_t stage)
        {
            switch (stage)
            {
                case GL_VERTEX_SHADER:    return ".cached_vulkan.vert";
                case GL_FRAGMENT_SHADER:  return ".cached_vulkan.frag";
            }
            OE_CORE_ASSERT(false, "No");
            return "";
        }


    }

    OpenGLShader::OpenGLShader(const std::string& filepath)
        : file_path(filepath)
    {
        OE_PROFILE_FUNCTION();

        Utils::CreateCacheDirectoryIfNeeded();

        std::string source = ReadFile(filepath);
        auto shader_sources = PreProcess(source);

        {
            CompileOrGetVulkanBinaries(shader_sources);
            CompileOrGetOpenGLBinaries();
            CreateProgram();
        }

        // Extract name from filepath
        auto last_slash = filepath.find_last_of("/\\");
        last_slash = last_slash == std::string::npos ? 0 : last_slash + 1;
        auto lastDot = filepath.rfind('.');
        auto count = lastDot == std::string::npos ? filepath.size() - last_slash : lastDot - last_slash;
        name = filepath.substr(last_slash, count);
    }

    OpenGLShader::OpenGLShader(const std::string& name, const std::string& vertex_src, const std::string& fragment_src)
        : name(name)
    {
        OE_PROFILE_FUNCTION();

        std::unordered_map<GLenum, std::string> sources;
        sources[GL_VERTEX_SHADER] = vertex_src;
        sources[GL_FRAGMENT_SHADER] = fragment_src;

        CompileOrGetVulkanBinaries(sources);
        CompileOrGetOpenGLBinaries();
        CreateProgram();
    }

    OpenGLShader::~OpenGLShader()
    {
        OE_PROFILE_FUNCTION();

        glDeleteProgram(id);
    }

    std::string OpenGLShader::ReadFile(const std::string& filepath)
    {
        OE_PROFILE_FUNCTION();

        std::string result;
        std::ifstream in(filepath, std::ios::in | std::ios::binary); // ifstream closes itself due to RAII
        if (in)
        {
            in.seekg(0, std::ios::end);
            size_t size = in.tellg();
            if (size != -1)
            {
                result.resize(size);
                in.seekg(0, std::ios::beg);
                in.read(&result[0], size);
            }
            else
            {
                OE_CORE_ERROR("Could not read from file '{0}'", filepath);
            }
        }
        else
        {
            OE_CORE_ERROR("Could not open file '{0}'", filepath);
        }

        return result;
    }

    std::unordered_map<GLenum, std::string> OpenGLShader::PreProcess(const std::string& source)
    {
        OE_PROFILE_FUNCTION();

        std::unordered_map<GLenum, std::string> shader_sources;

        const char* type_token = "#type";
        size_t type_token_length = strlen(type_token);
        size_t pos = source.find(type_token, 0); //Start of shader type declaration line
        while (pos != std::string::npos)
        {
            size_t eol = source.find_first_of("\r\n", pos); //End of shader type declaration line
            OE_CORE_ASSERT(eol != std::string::npos, "Syntax error");
            size_t begin = pos + type_token_length + 1; //Start of shader type name (after "#type " keyword)
            std::string type = source.substr(begin, eol - begin);
            OE_CORE_ASSERT(Utils::ShaderTypeFromString(type), "Invalid shader type specified");

            size_t next_line_pos = source.find_first_not_of("\r\n", eol); //Start of shader code after shader type declaration line
            OE_CORE_ASSERT(next_line_pos != std::string::npos, "Syntax error");
            pos = source.find(type_token, next_line_pos); //Start of next shader type declaration line

            shader_sources[Utils::ShaderTypeFromString(type)] = (pos == std::string::npos) ? source.substr(next_line_pos) : source.substr(next_line_pos, pos - next_line_pos);
        }

        return shader_sources;
    }

    void OpenGLShader::CompileOrGetVulkanBinaries(const std::unordered_map<GLenum, std::string>& shaderSources)
    {
        GLuint program = glCreateProgram();

        shaderc::Compiler compiler;
        shaderc::CompileOptions options;
        options.SetTargetEnvironment(shaderc_target_env_vulkan, shaderc_env_version_vulkan_1_2);
        const bool optimize = true;
        if (optimize)
            options.SetOptimizationLevel(shaderc_optimization_level_performance);

        std::filesystem::path cache_directory = Utils::GetCacheDirectory();

        auto& shader_data = vulkan_spirv;
        shader_data.clear();
        for (auto&& [stage, source] : shaderSources)
        {
            std::filesystem::path shader_file_path = file_path;
            std::filesystem::path cached_path = cache_directory / (shader_file_path.filename().string() + Utils::GLShaderStageCachedVulkanFileExtension(stage));

            std::ifstream in(cached_path, std::ios::in | std::ios::binary);
            if (in.is_open())
            {
                in.seekg(0, std::ios::end);
                auto size = in.tellg();
                in.seekg(0, std::ios::beg);

                auto& data = shader_data[stage];
                data.resize(size / sizeof(uint32_t));
                in.read((char*)data.data(), size);
            }
            else
            {
                shaderc::SpvCompilationResult module = compiler.CompileGlslToSpv(source, Utils::GLShaderStageToShaderC(stage), file_path.c_str(), options);
                if (module.GetCompilationStatus() != shaderc_compilation_status_success)
                {
                    OE_CORE_ERROR(module.GetErrorMessage());
                    OE_CORE_ASSERT(false, "No");
                }

                shader_data[stage] = std::vector<uint32_t>(module.cbegin(), module.cend());

                std::ofstream out(cached_path, std::ios::out | std::ios::binary);
                if (out.is_open())
                {
                    auto& data = shader_data[stage];
                    out.write((char*)data.data(), data.size() * sizeof(uint32_t));
                    out.flush();
                    out.close();
                }
            }
        }

        for (auto&& [stage, data]: shader_data)
            Reflect(stage, data);
    }

    void OpenGLShader::CompileOrGetOpenGLBinaries()
    {
        auto& shader_data = opengl_spirv;

        shaderc::Compiler compiler;
        shaderc::CompileOptions options;
        options.SetTargetEnvironment(shaderc_target_env_opengl, shaderc_env_version_opengl_4_5);
        const bool optimize = false;
        if (optimize)
            options.SetOptimizationLevel(shaderc_optimization_level_performance);

        std::filesystem::path cache_directory = Utils::GetCacheDirectory();

        shader_data.clear();
        m_OpenGLSourceCode.clear();
        for (auto&& [stage, spirv] : vulkan_spirv)
        {
            std::filesystem::path shader_file_path = file_path;
            std::filesystem::path cached_path = cache_directory / (shader_file_path.filename().string() + Utils::GLShaderStageCachedOpenGLFileExtension(stage));

            std::ifstream in(cached_path, std::ios::in | std::ios::binary);
            if (in.is_open())
            {
                in.seekg(0, std::ios::end);
                auto size = in.tellg();
                in.seekg(0, std::ios::beg);

                auto& data = shader_data[stage];
                data.resize(size / sizeof(uint32_t));
                in.read((char*)data.data(), size);
            }
            else
            {
                spirv_cross::CompilerGLSL glsl_compiler(spirv);
                m_OpenGLSourceCode[stage] = glsl_compiler.compile();
                auto& source = m_OpenGLSourceCode[stage];

                shaderc::SpvCompilationResult module = compiler.CompileGlslToSpv(source, Utils::GLShaderStageToShaderC(stage), file_path.c_str());
                if (module.GetCompilationStatus() != shaderc_compilation_status_success)
                {
                    OE_CORE_ERROR(module.GetErrorMessage());
                    OE_CORE_ASSERT(false, "No");
                }

                shader_data[stage] = std::vector<uint32_t>(module.cbegin(), module.cend());

                std::ofstream out(cached_path, std::ios::out | std::ios::binary);
                if (out.is_open())
                {
                    auto& data = shader_data[stage];
                    out.write((char*)data.data(), data.size() * sizeof(uint32_t));
                    out.flush();
                    out.close();
                }
            }
        }
    }

    void OpenGLShader::CreateProgram()
    {
        GLuint program = glCreateProgram();

        std::vector<GLuint> shader_ids;
        for (auto&& [stage, spirv] : opengl_spirv)
        {
            GLuint shaderID = shader_ids.emplace_back(glCreateShader(stage));
            glShaderBinary(1, &shaderID, GL_SHADER_BINARY_FORMAT_SPIR_V, spirv.data(), spirv.size() * sizeof(uint32_t));
            glSpecializeShader(shaderID, "main", 0, nullptr, nullptr);
            glAttachShader(program, shaderID);
        }

        glLinkProgram(program);

        GLint is_linked;
        glGetProgramiv(program, GL_LINK_STATUS, &is_linked);
        if (is_linked == GL_FALSE)
        {
            GLint max_length;
            glGetProgramiv(program, GL_INFO_LOG_LENGTH, &max_length);

            std::vector<GLchar> info_log(max_length);
            glGetProgramInfoLog(program, max_length, &max_length, info_log.data());
            OE_CORE_ERROR("Shader linking failed ({0}):\n{1}", file_path, info_log.data());

            glDeleteProgram(program);

            for (auto id : shader_ids)
                glDeleteShader(id);
        }

        for (auto id : shader_ids)
        {
            glDetachShader(program, id);
            glDeleteShader(id);
        }

        id = program;
    }

    void OpenGLShader::Reflect(GLenum stage, const std::vector<uint32_t>& shader_data)
    {
        spirv_cross::Compiler compiler(shader_data);
        spirv_cross::ShaderResources resources = compiler.get_shader_resources();

        OE_CORE_TRACE("OpenGLShader::Reflect - {0} {1}", Utils::GLShaderStageToString(stage), file_path);
        OE_CORE_TRACE("    {0} uniform buffers", resources.uniform_buffers.size());
        OE_CORE_TRACE("    {0} resources", resources.sampled_images.size());

        OE_CORE_TRACE("Uniform buffers:");
        for (const auto& resource : resources.uniform_buffers)
        {
            const auto& buffer_type = compiler.get_type(resource.base_type_id);
            uint32_t buffer_size = compiler.get_declared_struct_size(buffer_type);
            uint32_t binding = compiler.get_decoration(resource.id, spv::DecorationBinding);
            int member_count = buffer_type.member_types.size();

            OE_CORE_TRACE("  {0}", resource.name);
            OE_CORE_TRACE("    Size = {0}", buffer_size);
            OE_CORE_TRACE("    Binding = {0}", binding);
            OE_CORE_TRACE("    Members = {0}", member_count);
        }
    }

    void OpenGLShader::Bind() const
    { 
        OE_PROFILE_FUNCTION();
        glUseProgram(id); 
    }

    void OpenGLShader::Unbind() const
    {
        OE_PROFILE_FUNCTION();
        glUseProgram(0);
    }

    void OpenGLShader::SetBool(const std::string &name, bool value) const
    {
        OE_PROFILE_FUNCTION();

        UploadBool(name, value);
    }
    void OpenGLShader::SetInt(const std::string &name, int value) const
    {
        OE_PROFILE_FUNCTION();

        UploadInt(name, value);
    }
    void OpenGLShader::SetIntArray(const std::string &name, int* values, uint32_t count) const
    {
        OE_PROFILE_FUNCTION();

        UploadIntArray(name, values, count);
    }
    void OpenGLShader::SetFloat(const std::string &name, float value) const
    {
        OE_PROFILE_FUNCTION();

        UploadFloat(name, value);
    }
    void OpenGLShader::SetMat4(const std::string &name, const glm::mat4& value) const
    {
        OE_PROFILE_FUNCTION();

        UploadMat4(name, value);
    }
    void OpenGLShader::SetVec2(const std::string &name, const glm::vec2& value) const
    {
        OE_PROFILE_FUNCTION();

        UploadVec2(name, value);
    }
    void OpenGLShader::SetVec3(const std::string &name, const glm::vec3& value) const
    {
        OE_PROFILE_FUNCTION();

        UploadVec3(name, value);
    }
    void OpenGLShader::SetVec4(const std::string &name, const glm::vec4& value) const
    {
        OE_PROFILE_FUNCTION();

        UploadVec4(name, value);
    }

    void OpenGLShader::UploadBool(const std::string &name, bool value) const
    {         
        glUniform1i(glGetUniformLocation(id, name.c_str()), (int)value); 
    }
    void OpenGLShader::UploadInt(const std::string &name, int value) const
    { 
        glUniform1i(glGetUniformLocation(id, name.c_str()), value); 
    }
    void OpenGLShader::UploadIntArray(const std::string &name, int* values, uint32_t count) const
    { 
        glUniform1iv(glGetUniformLocation(id, name.c_str()), count, values); 
    }
    void OpenGLShader::UploadFloat(const std::string &name, float value) const
    { 
        glUniform1f(glGetUniformLocation(id, name.c_str()), value); 
    }
    void OpenGLShader::UploadMat4(const std::string &name, const glm::mat4& value) const
    {
        glUniformMatrix4fv(glGetUniformLocation(id, name.c_str()), 1, GL_FALSE, glm::value_ptr(value)); 
    }
    void OpenGLShader::UploadVec2(const std::string &name, const glm::vec2& value) const
    {
        glUniform2fv(glGetUniformLocation(id, name.c_str()), 1, glm::value_ptr(value)); 
    }
    void OpenGLShader::UploadVec3(const std::string &name, const glm::vec3& value) const
    {
        glUniform3fv(glGetUniformLocation(id, name.c_str()), 1, glm::value_ptr(value)); 
    }
    void OpenGLShader::UploadVec4(const std::string &name, const glm::vec4& value) const
    {
        glUniform4fv(glGetUniformLocation(id, name.c_str()), 1, glm::value_ptr(value)); 
    }
}
