//
// Created by sfd on 25-4-24.
//

#include "OpenGLShader.h"

#include <fstream>
#include <vector>
#include <array>
#include <filesystem>
#include <Hazel/Core/Log.h>
#include <Hazel/Debug/Instrumentor.h>
#include <shaderc/shaderc.hpp>

#include <spirv_cross.hpp>
#include <spirv_glsl.hpp>
#include <Hazel/Core/Timer.h>

#include "glm/gtc/type_ptr.hpp"
#include "spdlog/fmt/bundled/compile.h"

namespace Hazel
{
    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;

            HZ_CORE_ERROR("Unknown shader type!");
            return 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";
            }
            HZ_CORE_ERROR("Unknown shader stage!");
            return nullptr;
        }

        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;
            }
            HZ_CORE_ERROR("Unknown shader stage!");
            return (shaderc_shader_kind)0;
        }

        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";
            }
            HZ_CORE_ERROR("cannot find opengl shader cache");
            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";
            }
            HZ_CORE_ERROR("cannot find vulkan shader cache");
            return "";
        }

        static const char* GetCacheDirectory()
        {
            return "assets/cache/shader/opengl";
        }

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

    static GLenum ShaderTypeFromString(const std::string& type)
    {
        if (type == "vertex") return GL_VERTEX_SHADER;
        if (type == "fragment" || type == "pixel") return GL_FRAGMENT_SHADER;
        if (type == "geometry") return GL_GEOMETRY_SHADER;
        if (type == "compute") return GL_COMPUTE_SHADER;
        HZ_CORE_ERROR("Unknown Shader Type");
        return 0;
    }


    OpenGLShader::OpenGLShader(const std::string& filepath)
        : m_FilePath(filepath)
    {
        HZ_PROFILE_FUNCTION();

        Utils::CreateCacheDirectoryIfNeeded();


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

        {
            Timer timer;
            CompileOrGetVulkanBinaries(shaderSource);
            CompileOrGetOpenGLBinaries();
            CreateProgram();
            HZ_CORE_INFO("shader create took {0}ms",timer.ElapsedMillis());
        }

        const auto lastSlash = filepath.find_last_of("\\/");
        const auto lastDot = filepath.find_last_of(".");
        m_Name = filepath.substr(lastSlash + 1, lastDot - lastSlash - 1);

        /*
        std::string shaderSrc = ReadFile(filepath);
        auto shaderSources = PreProcess(shaderSrc);
        Compile(shaderSources);

        // example: /assets/shaders/Texture.glsl
        auto lastSlash = filepath.find_last_of("\\/");
        auto lastDot = filepath.find_last_of(".");
        m_Name = filepath.substr(lastSlash + 1, lastDot - lastSlash - 1);


        auto lastSlash = filepath.find_last_of("\\/");
        lastSlash = lastSlash == std::string::npos ? 0 : lastSlash + 1;
        auto lastDot = filepath.rfind(".");
        auto count = lastDot == std::string::npos ? filepath.size() - lastSlash : lastDot - lastSlash;
        m_Name = filepath.substr(lastSlash, count);
    */
    }

    OpenGLShader::OpenGLShader(const std::string& name, const std::string& vertexSrc, const std::string& fragmentSrc)
        : m_Name(name)
    {
        HZ_PROFILE_FUNCTION();

        std::unordered_map<GLenum, std::string> shaderSources;
        shaderSources[GL_VERTEX_SHADER] = vertexSrc;
        shaderSources[GL_FRAGMENT_SHADER] = fragmentSrc;
        // Compile(shaderSources);

        CompileOrGetVulkanBinaries(shaderSources);
        CompileOrGetOpenGLBinaries();
        CreateProgram();
    }


    OpenGLShader::~OpenGLShader()
    {
        HZ_PROFILE_FUNCTION();
        glDeleteProgram(m_RendererID);
    }

    void OpenGLShader::Bind() const
    {
        HZ_PROFILE_FUNCTION();
        glUseProgram(m_RendererID);
    }

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

    void OpenGLShader::SetFloat(const std::string& name, const float value)
    {
        HZ_PROFILE_FUNCTION();
        UploadUniformFloat(name, value);
    }

    void OpenGLShader::SetFloat3(const std::string& name, const glm::vec3& value)
    {
        HZ_PROFILE_FUNCTION();
        UploadUniformFloat3(name, value);
    }

    void OpenGLShader::SetFloat4(const std::string& name, const glm::vec4& value)
    {
        HZ_PROFILE_FUNCTION();
        UploadUniformFloat4(name, value);
    }

    void OpenGLShader::SetMat3(const std::string& name, const glm::mat3& value)
    {
        HZ_PROFILE_FUNCTION();
        UploadUniformMat3(name, value);
    }

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

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

    void OpenGLShader::SetIntv(const std::string& name, int* values, uint32_t count)
    {
        HZ_PROFILE_FUNCTION();
        UploadUniformIntv(name, values, count);
    }


    std::string OpenGLShader::ReadFile(const std::string& filepath)
    {
        HZ_PROFILE_FUNCTION();
        std::ifstream in(filepath, std::ios::in | std::ios::binary);
        // std::ifstream in(filepath, std::ios::in, std::ios::binary); // window platform api
        if (!in.is_open())
        {
            HZ_CORE_ERROR("Failed to open shader file {0}", filepath);
            return "";
        }

        std::string result;
        in.seekg(0, std::ios::end);
        result.resize(in.tellg());
        in.seekg(0, std::ios::beg);
        in.read(&result[0], result.size());
        in.close();

        return result;
    }

    std::unordered_map<GLenum, std::string> OpenGLShader::PreProcess(const std::string& source)
    {
        HZ_PROFILE_FUNCTION();
        std::unordered_map<GLenum, std::string> shaderSource;
        const char* typeToken = "#type";
        size_t typeTokenLength = strlen(typeToken);
        size_t pos = source.find(typeToken, 0);

        while (pos != std::string::npos)
        {
            size_t eol = source.find_first_of("\r\n", pos);
            if (eol == std::string::npos)
            {
                HZ_CORE_ERROR("[{0}:{1}]: Failed to parse shader source", __FILE__, __LINE__);
                return shaderSource;
            }
            size_t start = pos + typeTokenLength + 1;
            std::string type = source.substr(start, eol - start);
            if (!ShaderTypeFromString(type))
            {
                HZ_CORE_ERROR("[{0}:{1}]: Unsupported shader type", __FILE__, __LINE__);
                return shaderSource;
            }

            size_t nextLinePos = source.find_first_of("\r\n", eol);
            pos = source.find(typeToken, nextLinePos);
            shaderSource[ShaderTypeFromString(type)] = source.substr(nextLinePos, pos - (nextLinePos == std::string::npos ? source.size() - 1 : nextLinePos));

        }
        return shaderSource;

    }

    void OpenGLShader::Compile(std::unordered_map<GLenum, std::string>& shaderSource)
    {
        HZ_PROFILE_FUNCTION();
        GLuint program = glCreateProgram();
        // std::vector<GLenum> glShaderIDs;
        // glShaderIDs.reserve(shaderSource.size());

        if (shaderSource.size() > 2)
        {
            HZ_CORE_ERROR("Only support 2 shaders Now!");
        }
        std::array<GLenum, 2> glShaderIDs;
        int glShaderIDIndex = 0;
        for (auto& shaderSource : shaderSource)
        {
            GLenum type = shaderSource.first;
            const std::string source = shaderSource.second;

            GLuint shader = glCreateShader(type);

            const char* shaderSrc = source.c_str();
            glShaderSource(shader, 1, &shaderSrc, 0);

            glCompileShader(shader);

            GLint compileSuccess = 0;
            glGetShaderiv(shader, GL_COMPILE_STATUS, &compileSuccess);
            if (compileSuccess == GL_FALSE)
            {
                int length = 0;
                glGetShaderiv(shader, GL_INFO_LOG_LENGTH, &length);

                std::vector<GLchar> infoLog(length);
                glGetShaderInfoLog(shader, length, &length, &infoLog[0]);

                glDeleteShader(shader);
                HZ_CORE_ERROR("shader compilation failed!");
                HZ_CORE_ERROR("{0}", infoLog.data());
                return;
            }

            glAttachShader(program, shader);
            glShaderIDs[glShaderIDIndex++] = shader;
        }



        glLinkProgram(program);

        GLint compileSuccess = 0;
        glGetProgramiv(program, GL_LINK_STATUS, &compileSuccess);
        if (compileSuccess == GL_FALSE)
        {
            int length = 0;
            glGetProgramiv(program, GL_INFO_LOG_LENGTH, &length);
            std::vector<GLchar> infoLog(length);
            glGetProgramInfoLog(program, length, &length, &infoLog[0]);
            for (const auto id : glShaderIDs)
            {
                glDeleteShader(id);
            }

            HZ_CORE_ERROR("shaderProgram link failed!");
            HZ_CORE_ERROR("{0}", infoLog.data());
            return;
        }

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

        m_RendererID = program;
    }

    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 cacheDirectory = Utils::GetCacheDirectory();

        auto& shaderData = m_VulkanSPIRV;
        shaderData.clear();
        for (auto&& [stage, source] : shaderSources)
        {
            std::filesystem::path shaderFilePath = m_FilePath;
            std::filesystem::path cachedPath = cacheDirectory / (shaderFilePath.filename().string() + Utils::GLShaderStageCachedVulkanFileExtension(stage));

            std::ifstream in(cachedPath, 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 = shaderData[stage];
                data.resize(size / sizeof(uint32_t));
                in.read((char*)data.data(), size);
            }
            else
            {
                shaderc::SpvCompilationResult module = compiler.CompileGlslToSpv(source, Utils::GLShaderStageToShaderC(stage), m_FilePath.c_str(), options);
                if (module.GetCompilationStatus() != shaderc_compilation_status_success)
                {
                    HZ_CORE_ERROR(module.GetErrorMessage());
                    assert(-1);
                }

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

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

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

    void OpenGLShader::CompileOrGetOpenGLBinaries()
    {

        auto& shaderData = m_OpenGLSPIRV;

        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 cacheDirectory = Utils::GetCacheDirectory();

        shaderData.clear();
        m_OpenGLSourceCode.clear();
        for (auto&& [stage, spirv] : m_VulkanSPIRV)
        {
            std::filesystem::path shaderFilePath = m_FilePath;
            std::filesystem::path cachedPath = cacheDirectory / (shaderFilePath.filename().string() + Utils::GLShaderStageCachedOpenGLFileExtension(stage));

            std::ifstream in(cachedPath, 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 = shaderData[stage];
                data.resize(size / sizeof(uint32_t));
                in.read((char*)data.data(), size);
            }
            else
            {
                spirv_cross::CompilerGLSL glslCompiler(spirv);
                m_OpenGLSourceCode[stage] = glslCompiler.compile();
                auto& source = m_OpenGLSourceCode[stage];

                shaderc::SpvCompilationResult module = compiler.CompileGlslToSpv(source, Utils::GLShaderStageToShaderC(stage), m_FilePath.c_str());
                if (module.GetCompilationStatus() != shaderc_compilation_status_success)
                {
                    HZ_CORE_ERROR(module.GetErrorMessage());
                    assert(-1);
                }

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

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

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

        std::vector<GLuint> shaderIDs;
        for (auto&& [stage, spirv] : m_OpenGLSPIRV)
        {
            GLuint shaderID = shaderIDs.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 isLinked;
        glGetProgramiv(program, GL_LINK_STATUS, &isLinked);
        if (isLinked == GL_FALSE)
        {
            GLint maxLength;
            glGetProgramiv(program, GL_INFO_LOG_LENGTH, &maxLength);

            std::vector<GLchar> infoLog(maxLength);
            glGetProgramInfoLog(program, maxLength, &maxLength, infoLog.data());
            HZ_CORE_ERROR("Shader linking failed ({0}):\n{1}", m_FilePath, infoLog.data());

            glDeleteProgram(program);

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

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

        m_RendererID = program;
    }

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

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

		HZ_CORE_TRACE("Uniform buffers:");
		for (const auto& resource : resources.uniform_buffers)
		{
			const auto& bufferType = compiler.get_type(resource.base_type_id);
			uint32_t bufferSize = compiler.get_declared_struct_size(bufferType);
			uint32_t binding = compiler.get_decoration(resource.id, spv::DecorationBinding);
			int memberCount = bufferType.member_types.size();

			HZ_CORE_TRACE("  {0}", resource.name);
			HZ_CORE_TRACE("    Size = {0}", bufferSize);
			HZ_CORE_TRACE("    Binding = {0}", binding);
			HZ_CORE_TRACE("    Members = {0}", memberCount);
		}
    }

    void OpenGLShader::UploadUniformFloat3(const std::string& name, const glm::vec3& value) const
    {
        GLint location = glGetUniformLocation(m_RendererID, name.c_str());
        glUniform4fv(location, 1, glm::value_ptr(value));
    }

    void OpenGLShader::UploadUniformFloat4(const std::string& name, const glm::vec4& value) const
    {
        GLint location = glGetUniformLocation(m_RendererID, name.c_str());
        glUniform4fv(location, 1, glm::value_ptr(value));
    }

    void OpenGLShader::UploadUniformMat3(const std::string& name, const glm::mat3& matrix) const
    {
        GLint location = glGetUniformLocation(m_RendererID, name.c_str());
        glUniformMatrix3fv(location, 1, GL_FALSE, glm::value_ptr(matrix));
    }

    void OpenGLShader::UploadUniformMat4(const std::string& name, const glm::mat4& matrix) const
    {
        GLint location = glGetUniformLocation(m_RendererID, name.c_str());
        glUniformMatrix4fv(location, 1, GL_FALSE, glm::value_ptr(matrix));
    }

    void OpenGLShader::UploadUniformInt(const std::string& name, const int value) const
    {
        const GLint location = glGetUniformLocation(m_RendererID, name.c_str());
        glUniform1i(location, value);
    }

    void OpenGLShader::UploadUniformIntv(const std::string& name, int* values, uint32_t count) const
    {
        const GLint location = glGetUniformLocation(m_RendererID, name.c_str());
        glUniform1iv(location, count, values);
    }

    void OpenGLShader::UploadUniformFloat(const std::string& name, const float value) const
    {
        const GLint location = glGetUniformLocation(m_RendererID, name.c_str());
        glUniform1f(location, value);
    }
}
