#include "pch.h"

#include "ShaderManagement.h"

#include <spdlog/spdlog.h> 
#include <stdexcept>
#include <iostream> 
#include <glm/glm.hpp> 

#include "../../systems/RenderSystem.h"
#include "../../systems/sources/ShaderProgram.h"
#include "../../systems/sources/Shader.h"
#include "../../systems/sources/CompleteShader.h" 
#include "def/TextShader.h"
#include "def/SolidColorShader.h"
#include "def/Shader2D.h"

std::map<std::string, std::unique_ptr<ShaderProgram>> ShaderManagement::s_shaderMap;
ShaderProgram* ShaderManagement::s_defaultProgram = nullptr;

/**
 * @brief 使用 helper 函数绕过 std::initializer_list 的拷贝限制。
 * * * @return std::vector<std::unique_ptr<CompleteShader>>
 */
static std::vector<std::unique_ptr<CompleteShader>> create_shader_list()
{
    std::vector<std::unique_ptr<CompleteShader>> list;
    list.push_back(std::make_unique<Shader2D>());
    list.push_back(std::make_unique<SolidColorShader>());
    list.push_back(std::make_unique<TextShader>());
    return list;
}

std::vector<std::unique_ptr<CompleteShader>> ShaderManagement::s_shaderList = create_shader_list();

void ShaderManagement::compileAllShaders()
{
    RenderSystem::assertOnRenderThread();

    for (const auto& completeShader : s_shaderList) {
        compileShaderProgram(completeShader.get());
    }

    if (s_defaultProgram == nullptr && !s_shaderMap.empty()) {
        s_defaultProgram = s_shaderMap.begin()->second.get();
    }
}

void ShaderManagement::compileShaderProgram(const CompleteShader* completeShader)
{
    const std::string& shaderName = completeShader->getShaderName();

    try {
        const Shader& vertexShaderRef = completeShader->getVertexShader();
        GLuint vsId = compileShader(GL_VERTEX_SHADER, vertexShaderRef.getShaderCode(),
            vertexShaderRef.getShaderName());

        const Shader& fragmentShaderRef = completeShader->getFragmentShader();
        GLuint fsId = compileShader(GL_FRAGMENT_SHADER, fragmentShaderRef.getShaderCode(),
            fragmentShaderRef.getShaderName());

        GLuint programId = linkProgram(vsId, fsId, shaderName);
        std::unique_ptr<ShaderProgram> shaderProgram = std::make_unique<ShaderProgram>(programId);

        if (completeShader->isDefaultShader()) {
            s_defaultProgram = shaderProgram.get();
            setupDefaultUniforms(s_defaultProgram);
        }
        else {
            completeShader->setDefaultUniforms(shaderProgram.get());
        }

        s_shaderMap[shaderName] = std::move(shaderProgram);
        RenderSystem::deleteShader(vsId);
        RenderSystem::deleteShader(fsId);

        spdlog::info("成功编译着色器: {}", shaderName);

    }
    catch (const std::runtime_error& e) {
        spdlog::error("编译着色器失败: {}", shaderName);
        std::cerr << "Shader compilation failed: " << shaderName << "\nError: " << e.what() << std::endl;
        throw;
    }
}

void ShaderManagement::setupDefaultUniforms(ShaderProgram* program)
{
    if (!program) return;
    program->use();
    setUniformInt(program, "uTexture", 0);
    setUniformFloat(program, "uOpacity", 1.0f);
    setUniformVec4(program, "uColor", glm::vec4(1.0f, 1.0f, 1.0f, 1.0f));
    setUniformInt(program, "uBlendMode", 0);
    setUniformInt(program, "uDebugMode", 0);
    setUniformInt(program, "uLightCount", 0);
    program->stop();
    RenderSystem::checkGLError("setupDefaultUniforms");
}

GLuint ShaderManagement::compileShader(GLenum type, const std::string& source, const std::string& shaderName)
{
    GLuint shaderId = RenderSystem::createShader(type);
    RenderSystem::shaderSource(shaderId, source);
    RenderSystem::compileShader(shaderId);
    if (RenderSystem::getShaderi(shaderId, GL_COMPILE_STATUS) != GL_TRUE) {
        std::string log = RenderSystem::getShaderInfoLog(shaderId);
        RenderSystem::deleteShader(shaderId);
        throw std::runtime_error("着色器编译失败 [" + shaderName + "]:\n" + log);
    }
    return shaderId;
}

GLuint ShaderManagement::linkProgram(GLuint vertexShaderId, GLuint fragmentShaderId, const std::string& programName)
{
    GLuint programId = RenderSystem::createProgram();
    RenderSystem::attachShader(programId, vertexShaderId);
    RenderSystem::attachShader(programId, fragmentShaderId);
    RenderSystem::linkProgram(programId);
    if (RenderSystem::getProgrami(programId, GL_LINK_STATUS) != GL_TRUE) {
        std::string log = RenderSystem::getProgramInfoLog(programId);
        RenderSystem::deleteProgram(programId);
        throw std::runtime_error("着色器程序链接失败 [" + programName + "]:\n" + log);
    }
    validateProgram(programId, programName);
    return programId;
}

void ShaderManagement::validateProgram(GLuint programId, const std::string& programName)
{
    GLint validateStatus = RenderSystem::getProgrami(programId, GL_VALIDATE_STATUS);

    if (validateStatus != GL_TRUE) {
        std::string log = RenderSystem::getProgramInfoLog(programId);
        spdlog::warn("着色器程序验证警告 [{}]: {}", programName, log);
    }
}

ShaderProgram* ShaderManagement::getDefaultProgram()
{
    return s_defaultProgram;
}

ShaderProgram* ShaderManagement::getShaderProgram(const std::string& name)
{
    auto it = s_shaderMap.find(name);
    if (it != s_shaderMap.end()) {
        return it->second.get();
    }
    return nullptr;
}


void ShaderManagement::cleanup()
{
    RenderSystem::assertOnRenderThread();
    s_shaderMap.clear();
    s_shaderList.clear();
    s_defaultProgram = nullptr;
    spdlog::info("ShaderManagement cleanup completed.");
}

void ShaderManagement::setUniformInt(ShaderProgram* program, const std::string& name, GLint value)
{
    if (!program) return;
    program->use();
    GLint location = program->getUniformLocation(name);
    if (location != -1) {
        RenderSystem::uniform1i(location, value);
    }
}

void ShaderManagement::setUniformFloat(ShaderProgram* program, const std::string& name, GLfloat value)
{
    if (!program) return;
    program->use();
    GLint location = program->getUniformLocation(name);
    if (location != -1) {
        RenderSystem::uniform1f(location, value);
    }
}

void ShaderManagement::setUniformVec2(ShaderProgram* program, const std::string& name, GLfloat x, GLfloat y)
{
    if (!program) return;
    program->use();
    GLint location = program->getUniformLocation(name);
    if (location != -1) {
        RenderSystem::uniform2f(location, x, y);
    }
}

void ShaderManagement::setUniformVec2(ShaderProgram* program, const std::string& name, const glm::vec2& vec)
{
    if (!program) return;
    program->use();
    GLint location = program->getUniformLocation(name);
    if (location != -1) {
        RenderSystem::uniform2f(location, vec);
    }
}

void ShaderManagement::setUniformVec3(ShaderProgram* program, const std::string& name, GLfloat x, GLfloat y, GLfloat z)
{
    if (!program) return;
    program->use();
    GLint location = program->getUniformLocation(name);
    if (location != -1) {
        RenderSystem::uniform3f(location, x, y, z);
    }
}

void ShaderManagement::setUniformVec3(ShaderProgram* program, const std::string& name, const glm::vec3& vec)
{
    if (!program) return;
    program->use();
    GLint location = program->getUniformLocation(name);
    if (location != -1) {
        RenderSystem::uniform3f(location, vec);
    }
}

void ShaderManagement::setUniformVec4(ShaderProgram* program, const std::string& name, GLfloat x, GLfloat y, GLfloat z, GLfloat w)
{
    if (!program) return;
    program->use();
    GLint location = program->getUniformLocation(name);
    if (location != -1) {
        RenderSystem::uniform4f(location, x, y, z, w);
    }
}

void ShaderManagement::setUniformVec4(ShaderProgram* program, const std::string& name, const glm::vec4& vec)
{
    if (!program) return;
    program->use();
    GLint location = program->getUniformLocation(name);
    if (location != -1) {
        RenderSystem::uniform4f(location, vec);
    }
}

void ShaderManagement::setUniformMat3(ShaderProgram* program, const std::string& name, const glm::mat3& matrix)
{
    if (!program) return;
    program->use();
    GLint location = program->getUniformLocation(name);
    if (location != -1) {
        RenderSystem::uniformMatrix3(location, matrix, false);
    }
}