#include "Engine.h"
#include <iostream>
#include <glm/gtc/type_ptr.hpp>
#define STB_IMAGE_IMPLEMENTATION
#include "stb_image.h"
#include "Entity/Entity.h"

// Static member definitions.
GLFWwindow* Engine::window = nullptr;
GLuint Engine::framebuffer = 0;
GLuint Engine::colorTexture = 0;
GLuint Engine::depthRenderbuffer = 0;
GLuint Engine::shaderProgram = 0;
float Engine::rotationAngle = 0.0f;
int Engine::fbWidth = 640;
int Engine::fbHeight = 400;




// Global normal map texture (if needed for legacy models; otherwise each model handles its own)
GLuint normalMapTexture = 0;

bool Engine::Init() {
    if (!glfwInit()) {
        std::cout << "Failed to initialize GLFW\n";
        return false;
    }
    glfwWindowHint(GLFW_CONTEXT_VERSION_MAJOR, 3);
    glfwWindowHint(GLFW_CONTEXT_VERSION_MINOR, 3);
    glfwWindowHint(GLFW_OPENGL_PROFILE, GLFW_OPENGL_CORE_PROFILE);

    window = glfwCreateWindow(1280, 800, "Engine Window", nullptr, nullptr);
    if (!window) {
        std::cout << "Failed to create GLFW window\n";
        glfwTerminate();
        return false;
    }
    glfwMakeContextCurrent(window);

    glewExperimental = GL_TRUE;
    if (glewInit() != GLEW_OK) {
        std::cout << "Failed to initialize GLEW\n";
        return false;
    }

    int width, height;
    glfwGetFramebufferSize(window, &width, &height);
    glViewport(0, 0, width, height);

    // Create framebuffer.
    glGenFramebuffers(1, &framebuffer);
    ResizeFramebuffer(fbWidth, fbHeight);

    // Setup scene-wide shader (this shader is now used to render all models)
    if (!SetupScene()) {
        std::cout << "Failed to set up scene\n";
        return false;
    }
    return true;
}

GLFWwindow* Engine::GetWindow() {
    return window;
}

void Engine::ResizeFramebuffer(int width, int height) {
    // Avoid division by zero.
    if (height <= 0)
        height = 1;

    // Define the desired target aspect ratio (e.g., 16:9).
    const float targetAspect = 16.0f / 9.0f;
    float currentAspect = static_cast<float>(width) / static_cast<float>(height);
    
    // Adjust dimensions to maintain the target aspect ratio.
    int newWidth = width;
    int newHeight = height;
    if (currentAspect > targetAspect) {
        newWidth = static_cast<int>(height * targetAspect);
    } else if (currentAspect < targetAspect) {
        newHeight = static_cast<int>(width / targetAspect);
    }
    
    fbWidth = newWidth;
    fbHeight = newHeight;
    
    glBindFramebuffer(GL_FRAMEBUFFER, framebuffer);

    // Delete old attachments if they exist.
    if (colorTexture) {
        glDeleteTextures(1, &colorTexture);
    }
    if (depthRenderbuffer) {
        glDeleteRenderbuffers(1, &depthRenderbuffer);
    }

    // Create color texture using GL_RGBA.
    glGenTextures(1, &colorTexture);
    glBindTexture(GL_TEXTURE_2D, colorTexture);
    glTexImage2D(GL_TEXTURE_2D, 0, GL_RGBA, fbWidth, fbHeight, 0, GL_RGBA, GL_UNSIGNED_BYTE, NULL);
    glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, GL_LINEAR);
    glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, GL_LINEAR);
    glFramebufferTexture2D(GL_FRAMEBUFFER, GL_COLOR_ATTACHMENT0, GL_TEXTURE_2D, colorTexture, 0);

    // Create depth renderbuffer.
    glGenRenderbuffers(1, &depthRenderbuffer);
    glBindRenderbuffer(GL_RENDERBUFFER, depthRenderbuffer);
    glRenderbufferStorage(GL_RENDERBUFFER, GL_DEPTH24_STENCIL8, fbWidth, fbHeight);
    glFramebufferRenderbuffer(GL_FRAMEBUFFER, GL_DEPTH_STENCIL_ATTACHMENT, GL_RENDERBUFFER, depthRenderbuffer);

    GLenum status = glCheckFramebufferStatus(GL_FRAMEBUFFER);
    if (status != GL_FRAMEBUFFER_COMPLETE) {
        std::cout << "Framebuffer is not complete! Status: " << status << std::endl;
    }
    glBindFramebuffer(GL_FRAMEBUFFER, 0);
}

GLuint Engine::CompileShader(const char* vertexSrc, const char* fragmentSrc) {
    GLuint vertexShader = glCreateShader(GL_VERTEX_SHADER);
    glShaderSource(vertexShader, 1, &vertexSrc, nullptr);
    glCompileShader(vertexShader);
    int success;
    glGetShaderiv(vertexShader, GL_COMPILE_STATUS, &success);
    if (!success) {
        char infoLog[512];
        glGetShaderInfoLog(vertexShader, 512, nullptr, infoLog);
        std::cout << "Vertex shader compilation failed: " << infoLog << std::endl;
        return 0;
    }
    GLuint fragmentShader = glCreateShader(GL_FRAGMENT_SHADER);
    glShaderSource(fragmentShader, 1, &fragmentSrc, nullptr);
    glCompileShader(fragmentShader);
    glGetShaderiv(fragmentShader, GL_COMPILE_STATUS, &success);
    if (!success) {
        char infoLog[512];
        glGetShaderInfoLog(fragmentShader, 512, nullptr, infoLog);
        std::cout << "Fragment shader compilation failed: " << infoLog << std::endl;
        return 0;
    }
    GLuint program = glCreateProgram();
    glAttachShader(program, vertexShader);
    glAttachShader(program, fragmentShader);
    glLinkProgram(program);
    glGetProgramiv(program, GL_LINK_STATUS, &success);
    if (!success) {
        char infoLog[512];
        glGetProgramInfoLog(program, 512, nullptr, infoLog);
        std::cout << "Shader program linking failed: " << infoLog << std::endl;
        return 0;
    }
    glDeleteShader(vertexShader);
    glDeleteShader(fragmentShader);
    return program;
}

// SetupScene now creates the shader program used for all models.
// It no longer creates cube-specific VAOs, since ModelComponent will store mesh data.
bool Engine::SetupScene() {
    // Vertex shader: passes through vertex attributes.
    const char* vertexShaderSrc = R"(
        #version 330 core
        layout(location = 0) in vec3 aPos;
        layout(location = 1) in vec3 aNormal;
        layout(location = 2) in vec2 aTexCoords;
        layout(location = 3) in vec3 aTangent;
    
        uniform mat4 model;
        uniform mat4 view;
        uniform mat4 projection;
    
        out vec3 FragPos;
        out vec3 Normal;
        out vec2 TexCoords;
        out vec3 Tangent;
    
        void main() {
            FragPos = vec3(model * vec4(aPos, 1.0));
            Normal = mat3(transpose(inverse(model))) * aNormal;
            TexCoords = aTexCoords;
            Tangent = mat3(model) * aTangent;
            gl_Position = projection * view * vec4(FragPos, 1.0);
        }
    )";
    
    // Fragment shader: uses a normal map and material properties.
    const char* fragmentShaderSrc = R"(
        // Fragment shader:
#version 330 core
out vec4 FragColor;

in vec3 FragPos;
in vec3 Normal;
in vec2 TexCoords;
in vec3 Tangent;

uniform vec3 lightPositions[2];
uniform vec3 lightColors[2];
uniform int numLights;
uniform vec3 viewPos;
uniform sampler2D diffuseTexture;
uniform sampler2D normalMap;
uniform bool useNormalMap;  // NEW uniform to control normal mapping

// Material properties.
uniform vec3 materialDiffuse;
uniform vec3 materialSpecular;
uniform float materialShininess;

void main() {
    vec3 perturbedNormal = normalize(Normal);
    if(useNormalMap) {
        // Sample normal map.
        vec3 normMap = texture(normalMap, TexCoords).rgb;
        normMap = normalize(normMap * 2.0 - 1.0);
        // Flip Z if needed.
        normMap.z = -normMap.z;
        // Calculate tangent space basis.
        vec3 T = normalize(Tangent);
        vec3 B = normalize(cross(Normal, T));
        mat3 TBN = mat3(T, B, normalize(Normal));
        perturbedNormal = normalize(TBN * normMap);
    }
    
    vec3 diffuseTex = texture(diffuseTexture, TexCoords).rgb;

    vec3 ambient = 0.1 * materialDiffuse * diffuseTex;
    vec3 lighting = ambient;
    for(int i = 0; i < numLights; i++) {
        vec3 lightDir = normalize(lightPositions[i] - FragPos);
        float diff = max(dot(perturbedNormal, lightDir), 0.0);
        vec3 diffuse = diff * materialDiffuse * diffuseTex * lightColors[i];
        
        vec3 viewDir = normalize(viewPos - FragPos);
        vec3 reflectDir = reflect(-lightDir, perturbedNormal);
        float spec = pow(max(dot(viewDir, reflectDir), 0.0), materialShininess);
        vec3 specular = materialSpecular * spec * lightColors[i];
        
        lighting += diffuse + specular;
    }
    FragColor = vec4(lighting, 1.0);
}

    )";
    
    shaderProgram = CompileShader(vertexShaderSrc, fragmentShaderSrc);
    if (shaderProgram == 0) {
        return false;
    }
    
    return true;
}

// RenderScene now uses the new model system.
// For each entity of type CUBE with a valid ModelComponent, update uniforms and call Draw() on the model.
ImTextureID Engine::RenderScene(const glm::mat4 &view, const glm::mat4 &projection, const glm::vec3 &viewPos, const std::vector<Entity*>& entities) {
    glBindFramebuffer(GL_FRAMEBUFFER, framebuffer);
    glViewport(0, 0, fbWidth, fbHeight);
    glEnable(GL_DEPTH_TEST);
    glClearColor(0.1f, 0.1f, 0.1f, 1.0f);
    glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT);

    glUseProgram(shaderProgram);
    glUniformMatrix4fv(glGetUniformLocation(shaderProgram, "view"), 1, GL_FALSE, glm::value_ptr(view));
    glUniformMatrix4fv(glGetUniformLocation(shaderProgram, "projection"), 1, GL_FALSE, glm::value_ptr(projection));
    glUniform3f(glGetUniformLocation(shaderProgram, "viewPos"), viewPos.x, viewPos.y, viewPos.z);

    // Gather lights (up to 2) from entities of type LIGHT.
    glm::vec3 lightPositions[2] = { glm::vec3(0.0f), glm::vec3(0.0f) };
    glm::vec3 lightColors[2] = { glm::vec3(1.0f), glm::vec3(1.0f) };
    int lightCount = 0;
    for (auto e : entities) {
        if (e->GetType() == EntityType::LIGHT && lightCount < 2) {
            lightPositions[lightCount] = e->transform.position;
            if (e->lightComponent) {
                lightColors[lightCount] = e->lightComponent->color * e->lightComponent->intensity;
            }
            lightCount++;
        }
    }
    glUniform1i(glGetUniformLocation(shaderProgram, "numLights"), lightCount);
    if (lightCount > 0) {
        glUniform3fv(glGetUniformLocation(shaderProgram, "lightPositions"), lightCount, glm::value_ptr(lightPositions[0]));
        glUniform3fv(glGetUniformLocation(shaderProgram, "lightColors"), lightCount, glm::value_ptr(lightColors[0]));
    }

    // Render each cube entity using its ModelComponent.
for (auto e : entities) {
    if (e->GetType() == EntityType::CUBE && e->modelComponent) {
        glm::mat4 modelMatrix = e->transform.GetMatrix();
        glUniformMatrix4fv(glGetUniformLocation(shaderProgram, "model"), 1, GL_FALSE, glm::value_ptr(modelMatrix));

        // Loop through all submeshes in the model component.
        for (const auto &mesh : e->modelComponent->meshes) {
            // Set material properties for the current submesh.
            glUniform3fv(glGetUniformLocation(shaderProgram, "materialDiffuse"), 1, glm::value_ptr(mesh.diffuseColor));
            glUniform3fv(glGetUniformLocation(shaderProgram, "materialSpecular"), 1, glm::value_ptr(mesh.specularColor));
            glUniform1f(glGetUniformLocation(shaderProgram, "materialShininess"), mesh.shininess);

            // Bind the diffuse texture.
            glActiveTexture(GL_TEXTURE0);
            glBindTexture(GL_TEXTURE_2D, mesh.diffuseTexture);
            glUniform1i(glGetUniformLocation(shaderProgram, "diffuseTexture"), 0);

            // If you have a normal texture, bind it similarly (adjust as needed).
            // glActiveTexture(GL_TEXTURE1);
            // glBindTexture(GL_TEXTURE_2D, mesh.normalTexture);
            // glUniform1i(glGetUniformLocation(shaderProgram, "normalMap"), 1);

            // Bind the submesh's VAO and draw its elements.
            glBindVertexArray(mesh.VAO);
            glDrawElements(GL_TRIANGLES, static_cast<GLsizei>(mesh.indices.size()), GL_UNSIGNED_INT, 0);
            glBindVertexArray(0);
        }
    }
}


    glBindFramebuffer(GL_FRAMEBUFFER, 0);
    return (ImTextureID)(intptr_t)colorTexture;
}

ImTextureID Engine::GetFinalRenderingTexture() {
    return (ImTextureID)(intptr_t)colorTexture;
}

void Engine::Shutdown() {
    glDeleteProgram(shaderProgram);
    glDeleteFramebuffers(1, &framebuffer);
    glDeleteTextures(1, &colorTexture);
    glDeleteTextures(1, &normalMapTexture);
    glDeleteRenderbuffers(1, &depthRenderbuffer);
    glfwDestroyWindow(window);
    glfwTerminate();
}