#include "OpenGLShader.h"

#include <glad/glad.h>
#include <glm/gtc/type_ptr.hpp>

#include <fstream>
#include <iostream>

namespace Donut
{
    static uint32_t ShaderTypeFromString(const std::string& type)
    {
        if (type == "vertex")
            return GL_VERTEX_SHADER;
        if (type == "fragment" || type == "pixel")
            return GL_FRAGMENT_SHADER;
        if (type == "compute")
            return GL_COMPUTE_SHADER;
        return 0;
    }

    // Shaders are authored in Slang and compiled to Assets/Shaders/generated/
    // <name>.glsl by Tools/compile-shaders.sh. Given a legacy ".../<name>.glsl"
    // path, prefer that generated file when present; otherwise fall back to the
    // hand-written GLSL (e.g. shaders not yet ported to Slang).
    static std::string ResolveShaderPath(const std::string& filepath)
    {
        size_t slash = filepath.find_last_of("/\\");
        std::string dir  = (slash == std::string::npos) ? std::string() : filepath.substr(0, slash + 1);
        std::string file = (slash == std::string::npos) ? filepath : filepath.substr(slash + 1);
        size_t dot = file.rfind('.');
        std::string base = (dot == std::string::npos) ? file : file.substr(0, dot);

        std::string generated = dir + "generated/" + base + ".glsl";
        std::ifstream test(generated);
        if (test.good())
            return generated;
        return filepath;
    }

    OpenGLShader::OpenGLShader(const std::string& filepath)
    {
        std::string resolved = ResolveShaderPath(filepath);
        m_IsSlang = (resolved != filepath);
        std::string source = ReadFile(resolved);
        auto shaderSources = PreProcess(source);
        Compile(shaderSources);

        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)
    {
        std::unordered_map<uint32_t, std::string> sources;
        sources[GL_VERTEX_SHADER] = vertexSrc;
        sources[GL_FRAGMENT_SHADER] = fragmentSrc;
        Compile(sources);
    }

    OpenGLShader::OpenGLShader(const std::string& name, const std::string& computeSrc)
        : m_Name(name)
    {
        std::unordered_map<uint32_t, std::string> sources;
        sources[GL_COMPUTE_SHADER] = computeSrc;
        Compile(sources);
    }

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

    std::string OpenGLShader::ReadFile(const std::string& filepath)
    {
        std::string result;
        std::ifstream in(filepath, std::ios::in |
                                   std::ios::binary);

        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);
            }
        }
        return result;
    }

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

        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);
            size_t begin = pos + typeTokenLength + 1;
            std::string type = source.substr(begin, eol - begin);

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

        return shaderSources;
    }

    void OpenGLShader::Compile(const std::unordered_map<uint32_t, std::string>& shaderSources)
    {
        uint32_t program = glCreateProgram();
        std::vector<uint32_t> glShaderIDs(shaderSources.size());
        for (auto& kv : shaderSources)
        {
            uint32_t type = kv.first;
            const std::string& source = kv.second;

            uint32_t shader = glCreateShader(type);
            const char* sourceCStr = source.c_str();
            glShaderSource(shader, 1, &sourceCStr, 0);
            glCompileShader(shader);

            int isCompiled = 0;
            glGetShaderiv(shader, GL_COMPILE_STATUS, &isCompiled);
            if (isCompiled == GL_FALSE)
            {
                int maxLength = 0;
                glGetShaderiv(shader, GL_INFO_LOG_LENGTH, &maxLength);
                std::vector<char> infoLog(maxLength);
                glGetShaderInfoLog(shader, maxLength, &maxLength, &infoLog[0]);
                glDeleteShader(shader);
                for (auto id : glShaderIDs)
                    glDeleteShader(id);
                glDeleteProgram(program);
                m_RendererID = 0;
                // infoLog.data() is null when the driver returns an empty log
                // (e.g. macOS rejecting a compute shader); streaming a null
                // char* into std::cout calls strlen(NULL) and crashes.
                const char* log = infoLog.empty() ? "" : infoLog.data();
                std::cout << "Shader compilation failure!" << std::endl << log << std::endl;
                return;
            }
            glAttachShader(program, shader);
            glShaderIDs.push_back(shader);
        }

        m_RendererID = program;
        glLinkProgram(m_RendererID);

        int isLinked = 0;
        glGetProgramiv(m_RendererID, GL_LINK_STATUS, (int*)&isLinked);
        if (isLinked == GL_FALSE)
        {
            int maxLength = 0;
            glGetProgramiv(m_RendererID, GL_INFO_LOG_LENGTH, &maxLength);
            std::vector<char> infoLog(maxLength);
            glGetProgramInfoLog(m_RendererID, maxLength, &maxLength, &infoLog[0]);
            glDeleteProgram(m_RendererID);
            for (auto id : glShaderIDs)
                glDeleteShader(id);
            m_RendererID = 0;
            const char* log = infoLog.empty() ? "" : infoLog.data();
            std::cout << "Shader link failure!" << std::endl << log << std::endl;
            return;
        }

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

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

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

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

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

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

    void OpenGLShader::SetFloat2(const std::string& name, const glm::vec2& value)
    {
        UploadUniformFloat2(name, value);
    }

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

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

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

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

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

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

    void OpenGLShader::UploadUniformFloat2(const std::string& name, const glm::vec2& value)
    {
        int location = glGetUniformLocation(m_RendererID, name.c_str());
        glUniform2f(location, value.x, value.y);
    }

    void OpenGLShader::UploadUniformFloat3(const std::string& name, const glm::vec3& value)
    {
        int location = glGetUniformLocation(m_RendererID, name.c_str());
        glUniform3f(location, value.x, value.y, value.z);
    }

    void OpenGLShader::UploadUniformFloat4(const std::string& name, const glm::vec4& value)
    {
        int location = glGetUniformLocation(m_RendererID, name.c_str());
        glUniform4f(location, value.x, value.y, value.z, value.w);
    }

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

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

    void OpenGLShader::Dispatch(uint32_t x, uint32_t y, uint32_t z)
    {
        // Compute shaders require OpenGL 4.3+. On drivers that cap out earlier
        // (e.g. macOS, which is frozen at 4.1) glDispatchCompute is never
        // loaded and the pointer is null. Guard so we no-op instead of crash.
        if (m_RendererID == 0 || glDispatchCompute == nullptr)
            return;
        glDispatchCompute(x, y, z);
    }

    void OpenGLShader::DispatchIndirect(uint32_t offset)
    {
        if (m_RendererID == 0 || glDispatchComputeIndirect == nullptr)
            return;
        glDispatchComputeIndirect(offset);
    }

    void OpenGLShader::MemoryBarrier(uint32_t barriers)
    {
        if (glMemoryBarrier == nullptr)
            return;
        glMemoryBarrier(barriers);
    }
};