#include "ModelComponent.h"
#include "stb_image.h"
#define TINYOBJLOADER_IMPLEMENTATION
#include "tiny_obj_loader.h"

#include <iostream>
#include <sstream>
#include <map>
#include <algorithm>

// Constructor: initialize default values.
ModelComponent::ModelComponent()
    : modelPath("./assets/models/sponza.obj"),
      diffuseColor(1.0f, 1.0f, 1.0f),
      specularColor(1.0f, 1.0f, 1.0f),
      shininess(32.0f)
{
    // Load the model on construction.
    LoadModel(modelPath);
}

ModelComponent::~ModelComponent() {
    // Clean up each submesh.
    for (auto &mesh : meshes) {
        if(mesh.VAO) glDeleteVertexArrays(1, &mesh.VAO);
        if(mesh.VBO) glDeleteBuffers(1, &mesh.VBO);
        if(mesh.EBO) glDeleteBuffers(1, &mesh.EBO);
        if(mesh.diffuseTexture) glDeleteTextures(1, &mesh.diffuseTexture);
    }
}

void ModelComponent::SetupMesh(SubMesh &mesh) {
    glGenVertexArrays(1, &mesh.VAO);
    glGenBuffers(1, &mesh.VBO);
    glGenBuffers(1, &mesh.EBO);
    
    glBindVertexArray(mesh.VAO);
    
    glBindBuffer(GL_ARRAY_BUFFER, mesh.VBO);
    glBufferData(GL_ARRAY_BUFFER, mesh.vertices.size() * sizeof(Vertex), &mesh.vertices[0], GL_STATIC_DRAW);
    
    glBindBuffer(GL_ELEMENT_ARRAY_BUFFER, mesh.EBO);
    glBufferData(GL_ELEMENT_ARRAY_BUFFER, mesh.indices.size() * sizeof(GLuint), &mesh.indices[0], GL_STATIC_DRAW);
    
    // Vertex attributes:
    // Positions.
    glVertexAttribPointer(0, 3, GL_FLOAT, GL_FALSE, sizeof(Vertex), (void*)0);
    glEnableVertexAttribArray(0);
    // Normals.
    glVertexAttribPointer(1, 3, GL_FLOAT, GL_FALSE, sizeof(Vertex), (void*)offsetof(Vertex, Normal));
    glEnableVertexAttribArray(1);
    // TexCoords.
    glVertexAttribPointer(2, 2, GL_FLOAT, GL_FALSE, sizeof(Vertex), (void*)offsetof(Vertex, TexCoords));
    glEnableVertexAttribArray(2);
    // Tangents.
    glVertexAttribPointer(3, 3, GL_FLOAT, GL_FALSE, sizeof(Vertex), (void*)offsetof(Vertex, Tangent));
    glEnableVertexAttribArray(3);
    
    glBindVertexArray(0);
}

bool ModelComponent::LoadDiffuseTexture(const std::string &filepath, GLuint &textureID) {
    int w, h, channels;
    unsigned char* data = stbi_load(filepath.c_str(), &w, &h, &channels, 0);
    if (!data) {
        std::cout << "Failed to load diffuse texture: " << filepath << std::endl;
        return false;
    }
    glGenTextures(1, &textureID);
    glBindTexture(GL_TEXTURE_2D, textureID);
    GLenum format = (channels == 3) ? GL_RGB : GL_RGBA;
    glTexImage2D(GL_TEXTURE_2D, 0, format, w, h, 0, format, GL_UNSIGNED_BYTE, data);
    glGenerateMipmap(GL_TEXTURE_2D);
    glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_S, GL_REPEAT);
    glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_T, GL_REPEAT);
    glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, GL_LINEAR_MIPMAP_LINEAR);
    glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, GL_LINEAR);
    stbi_image_free(data);
    return true;
}

bool ModelComponent::LoadNormalTexture(const std::string &filepath) {
    // Similar to LoadDiffuseTexture but for normalTexture; not used in this submesh example.
    return true;
}

// Loads a model using tiny_obj_loader and groups faces by material.
bool ModelComponent::LoadModel(const std::string &filepath) {
    modelPath = filepath;
    tinyobj::attrib_t attrib;
    std::vector<tinyobj::shape_t> shapes;
    std::vector<tinyobj::material_t> materials;
    std::string warn, err;

    // Get the directory of the OBJ file.
    std::string baseDir = "";
    size_t lastSlash = filepath.find_last_of("/\\");
    if (lastSlash != std::string::npos) {
        baseDir = filepath.substr(0, lastSlash + 1);
    }

    bool ret = tinyobj::LoadObj(&attrib, &shapes, &materials, &warn, &err, filepath.c_str(), baseDir.c_str());
    if (!warn.empty()) {
        std::cout << "[Warning] " << warn << std::endl;
    }
    if (!err.empty()) {
        std::cerr << "[Error] " << err << std::endl;
    }
    if (!ret) {
        std::cerr << "[Error] Failed to load/parse OBJ file: " << filepath << std::endl;
        return false;
    }

    // Clear any existing submeshes.
    meshes.clear();

    // For each shape in the OBJ file.
    for (size_t s = 0; s < shapes.size(); s++) {
        // Use a map to group faces by material ID.
        std::map<int, SubMesh> submeshMap;
        size_t index_offset = 0;
        // Iterate over each face.
        for (size_t f = 0; f < shapes[s].mesh.num_face_vertices.size(); f++) {
            std::cout << "[Info] AssetManager: Loading Mesh: "<< i << std::endl;
            int fv = shapes[s].mesh.num_face_vertices[f];
            // Get material ID for this face.
            int matID = -1;
            if (shapes[s].mesh.material_ids.size() > f)
                matID = shapes[s].mesh.material_ids[f];
            // If submesh for this material does not exist, create it.
            if (submeshMap.find(matID) == submeshMap.end())
                submeshMap[matID] = SubMesh();
            
            SubMesh &currentMesh = submeshMap[matID];
            // Process each vertex in the face.
            for (size_t v = 0; v < fv; v++) {
                tinyobj::index_t idx = shapes[s].mesh.indices[index_offset + v];
                Vertex vertex;
                // Position.
                vertex.Position = glm::vec3(
                    attrib.vertices[3 * idx.vertex_index + 0],
                    attrib.vertices[3 * idx.vertex_index + 1],
                    attrib.vertices[3 * idx.vertex_index + 2]
                );
                // Normal.
                if (idx.normal_index >= 0 && 3 * idx.normal_index + 2 < attrib.normals.size()) {
                    vertex.Normal = glm::vec3(
                        attrib.normals[3 * idx.normal_index + 0],
                        attrib.normals[3 * idx.normal_index + 1],
                        attrib.normals[3 * idx.normal_index + 2]
                    );
                } else {
                    vertex.Normal = glm::vec3(0.0f);
                }
                // Texture Coordinates.
                if (idx.texcoord_index >= 0 && 2 * idx.texcoord_index + 1 < attrib.texcoords.size()) {
                    vertex.TexCoords = glm::vec2(
                        attrib.texcoords[2 * idx.texcoord_index + 0],
                        attrib.texcoords[2 * idx.texcoord_index + 1]
                    );
                } else {
                    vertex.TexCoords = glm::vec2(0.0f);
                }
                // Tangent not provided.
                vertex.Tangent = glm::vec3(0.0f);
                
                // Add vertex and index.
                currentMesh.vertices.push_back(vertex);
                currentMesh.indices.push_back(static_cast<GLuint>(currentMesh.vertices.size() - 1));
            }
            index_offset += fv;
        }

        // For each group (submesh) in this shape, set material properties and create buffers.
        for (auto &pair : submeshMap) {
            int matID = pair.first;
            SubMesh &subMesh = pair.second;
            // If there is a valid material, assign its properties.
            if (matID >= 0 && matID < static_cast<int>(materials.size())) {
                tinyobj::material_t mat = materials[matID];
                subMesh.diffuseColor = glm::vec3(mat.diffuse[0], mat.diffuse[1], mat.diffuse[2]);
                subMesh.specularColor = glm::vec3(mat.specular[0], mat.specular[1], mat.specular[2]);
                subMesh.shininess = mat.shininess;
                if (!mat.diffuse_texname.empty()) {
                    std::string texturePath = baseDir + mat.diffuse_texname;
                    LoadDiffuseTexture(texturePath, subMesh.diffuseTexture);
                }
            }
            // Setup the OpenGL buffers for this submesh.
            SetupMesh(subMesh);
            // Add submesh to our model.
            meshes.push_back(subMesh);
        }
    }

    return true;
}

// Draws the entire model by drawing each submesh.
void ModelComponent::Draw() {
    for (auto &mesh : meshes) {
        // Bind the appropriate texture if available.
        if(mesh.diffuseTexture) {
            glActiveTexture(GL_TEXTURE0);
            glBindTexture(GL_TEXTURE_2D, mesh.diffuseTexture);
        }
        glBindVertexArray(mesh.VAO);
        glDrawElements(GL_TRIANGLES, static_cast<GLsizei>(mesh.indices.size()), GL_UNSIGNED_INT, 0);
        glBindVertexArray(0);
    }
}