// src/Shader.cpp

#include "Shader.h"
#include <vector>  // Add this include

// Helper to read a file into a string
std::optional<std::string> Shader::ReadFile(const std::filesystem::path& path) {
    std::ifstream file(path, std::ios::in | std::ios::binary);
    if (!file) return std::nullopt;
    std::ostringstream contents;
    contents << file.rdbuf();
    return contents.str();
}

// Constructor (vertex + fragment)
Shader::Shader(const std::filesystem::path& vertex_path,
               const std::filesystem::path& fragment_path) {
    auto vert_src_opt = ReadFile(vertex_path);
    auto frag_src_opt = ReadFile(fragment_path);
    if (!vert_src_opt || !frag_src_opt) {
        std::cerr << "Failed to read shader files\n";
        return;
    }
    auto vert_res = TryCompileShader(GL_VERTEX_SHADER, *vert_src_opt, vertex_path.string());
    auto frag_res = TryCompileShader(GL_FRAGMENT_SHADER, *frag_src_opt, fragment_path.string());
    if (!vert_res.success || !frag_res.success) {
        std::cerr << vert_res.error_message << "\n" << frag_res.error_message << "\n";
        return;
    }
    std::vector<GLuint> shaders = { vert_res.shader_id, frag_res.shader_id };
    program_id_ = CreateProgram(shaders);
    // Cleanup shader objects
    glDeleteShader(vert_res.shader_id);
    glDeleteShader(frag_res.shader_id);
}

// Constructor (vertex + geometry + fragment)
Shader::Shader(const std::filesystem::path& vertex_path,
               const std::filesystem::path& geometry_path,
               const std::filesystem::path& fragment_path) {
    auto vert_src_opt = ReadFile(vertex_path);
    auto geom_src_opt = ReadFile(geometry_path);
    auto frag_src_opt = ReadFile(fragment_path);
    if (!vert_src_opt || !geom_src_opt || !frag_src_opt) {
        std::cerr << "Failed to read shader files\n";
        return;
    }
    auto vert_res = TryCompileShader(GL_VERTEX_SHADER, *vert_src_opt, vertex_path.string());
    auto geom_res = TryCompileShader(GL_GEOMETRY_SHADER, *geom_src_opt, geometry_path.string());
    auto frag_res = TryCompileShader(GL_FRAGMENT_SHADER, *frag_src_opt, fragment_path.string());
    if (!vert_res.success || !geom_res.success || !frag_res.success) {
        std::cerr << vert_res.error_message << "\n"
                  << geom_res.error_message << "\n"
                  << frag_res.error_message << "\n";
        return;
    }
    std::vector<GLuint> shaders = { vert_res.shader_id, geom_res.shader_id, frag_res.shader_id };
    program_id_ = CreateProgram(shaders);
    glDeleteShader(vert_res.shader_id);
    glDeleteShader(geom_res.shader_id);
    glDeleteShader(frag_res.shader_id);
}

Shader::~Shader() {
    Cleanup();
}

Shader::Shader(Shader&& other) noexcept {
    program_id_ = other.program_id_;
    other.program_id_ = 0;
}

Shader& Shader::operator=(Shader&& other) noexcept {
    if (this != &other) {
        Cleanup();
        program_id_ = other.program_id_;
        other.program_id_ = 0;
    }
    return *this;
}

void Shader::Bind() const {
    if (program_id_) glUseProgram(program_id_);
}

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

void Shader::Cleanup() {
    if (program_id_) {
        glDeleteProgram(program_id_);
        program_id_ = 0;
    }
}

GLuint Shader::CompileShader(GLenum type, const std::string& source) {
    GLuint id = glCreateShader(type);
    const char* src = source.c_str();
    glShaderSource(id, 1, &src, nullptr);
    glCompileShader(id);
    return id;
}

Shader::CompileResult Shader::TryCompileShader(GLenum type,
                                               const std::string& source,
                                               const std::string& name) {
    CompileResult result;
    GLuint shader = CompileShader(type, source);
    GLint status;
    glGetShaderiv(shader, GL_COMPILE_STATUS, &status);
    if (status == GL_FALSE) {
        GLint length;
        glGetShaderiv(shader, GL_INFO_LOG_LENGTH, &length);
        std::string log(length, ' ');
        glGetShaderInfoLog(shader, length, &length, log.data());
        result.success = false;
        result.error_message = "Error compiling shader " + name + ": " + log;
        glDeleteShader(shader);
    } else {
        result.success = true;
        result.shader_id = shader;
    }
    return result;
}

GLuint Shader::CreateProgram(const std::vector<GLuint>& shaders) {
    GLuint program = glCreateProgram();
    for (auto s : shaders) {
        glAttachShader(program, s);
    }
    glLinkProgram(program);
    GLint status;
    glGetProgramiv(program, GL_LINK_STATUS, &status);
    if (status == GL_FALSE) {
        GLint length;
        glGetProgramiv(program, GL_INFO_LOG_LENGTH, &length);
        std::string log(length, ' ');
        glGetProgramInfoLog(program, length, &length, log.data());
        std::cerr << "Error linking program: " << log << "\n";
        glDeleteProgram(program);
        return 0;
    }
    return program;
}

int Shader::GetUniformLocation(const std::string& name) {
    if (uniform_location_cache_.count(name)) {
        return uniform_location_cache_[name];
    }
    int location = glGetUniformLocation(program_id_, name.c_str());
    uniform_location_cache_[name] = location;
    return location;
}

void Shader::PrintActiveUniforms() const {
    GLint count;
    glGetProgramiv(program_id_, GL_ACTIVE_UNIFORMS, &count);
    for (GLint i = 0; i < count; ++i) {
        char name[256];
        GLsizei length;
        GLint size;
        GLenum type;
        glGetActiveUniform(program_id_, i, sizeof(name), &length, &size, &type, name);
        std::cout << "Uniform #" << i << ": " << name << "\n";
    }
}

void Shader::PrintActiveAttributes() const {
    GLint count;
    glGetProgramiv(program_id_, GL_ACTIVE_ATTRIBUTES, &count);
    for (GLint i = 0; i < count; ++i) {
        char name[256];
        GLsizei length;
        GLint size;
        GLenum type;
        glGetActiveAttrib(program_id_, i, sizeof(name), &length, &size, &type, name);
        std::cout << "Attribute #" << i << ": " << name << "\n";
    }
}

void Shader::SetUniform(const std::string& name, const glm::vec2& value) {
    glUniform2fv(GetUniformLocation(name), 1, &value.x);
}

void Shader::SetUniform(const std::string& name, const glm::vec3& value) {
    glUniform3fv(GetUniformLocation(name), 1, &value.x);
}

void Shader::SetUniform(const std::string& name, const glm::vec4& value) {
    glUniform4fv(GetUniformLocation(name), 1, &value.x);
}

void Shader::SetUniform(const std::string& name, const glm::mat3& value) {
    glUniformMatrix3fv(GetUniformLocation(name), 1, GL_FALSE, &value[0][0]);
}

void Shader::SetUniform(const std::string& name, const glm::mat4& value) {
    glUniformMatrix4fv(GetUniformLocation(name), 1, GL_FALSE, &value[0][0]);
}
