#include "simulation_state.h"
#include "rendering/renderer.h"
#include "rendering/vertex_array.h"
#include "rendering/vertex_buffer.h"
#include "rendering/index_buffer.h"
#include "rendering/shader.h"
#include "rendering/texture.h"

#include <iostream>
#include <glm/glm.hpp>
#include <glm/gtc/matrix_transform.hpp>
#include <GLFW/glfw3.h>

namespace Donut
{
    void SimulationState::on_enter()
    {
        m_camera = std::make_unique<Camera>(45.0f, 1280.0f / 720.0f, 0.1f, 100.0f);
        m_camera->set_position({ 0.0f, 0.0f, 3.0f });
        m_camera->set_mouse_sensitivity(0.1f);
        m_camera->set_movement_speed(5.0f);

        float vertices[] = 
        {
            -0.5f, -0.5f, -0.5f,  1.0f, 0.0f, 0.0f, 1.0f,  0.0f, 0.0f,
             0.5f, -0.5f, -0.5f,  0.0f, 1.0f, 0.0f, 1.0f,  1.0f, 0.0f,
             0.5f,  0.5f, -0.5f,  0.0f, 0.0f, 1.0f, 1.0f,  1.0f, 1.0f,
            -0.5f,  0.5f, -0.5f,  1.0f, 1.0f, 0.0f, 1.0f,  0.0f, 1.0f,
            -0.5f, -0.5f,  0.5f,  1.0f, 0.0f, 1.0f, 1.0f,  0.0f, 0.0f,
             0.5f, -0.5f,  0.5f,  0.0f, 1.0f, 1.0f, 1.0f,  1.0f, 0.0f,
             0.5f,  0.5f,  0.5f,  0.5f, 0.5f, 0.5f, 1.0f,  1.0f, 1.0f,
            -0.5f,  0.5f,  0.5f,  1.0f, 1.0f, 1.0f, 1.0f,  0.0f, 1.0f
        };

        uint32_t indices[] = 
        {
            0, 1, 2,  2, 3, 0,
            1, 5, 6,  6, 2, 1,
            5, 4, 7,  7, 6, 5,
            4, 0, 3,  3, 7, 4,
            3, 2, 6,  6, 7, 3,
            4, 5, 1,  1, 0, 4
        };

        m_vertex_array = std::shared_ptr<VertexArray>(VertexArray::create());
        m_vertex_buffer = std::shared_ptr<VertexBuffer>(VertexBuffer::create(vertices, sizeof(vertices)));
        
        VertexBufferLayout layout;
        layout.push<float>(3);
        layout.push<float>(4);
        layout.push<float>(2);
        m_vertex_buffer->set_layout(layout);

        m_vertex_array->add_vertex_buffer(m_vertex_buffer);

        m_index_buffer = std::shared_ptr<IndexBuffer>(IndexBuffer::create(indices, 36));
        m_vertex_array->set_index_buffer(m_index_buffer);

        m_shader = std::shared_ptr<Shader>(Shader::create("assets/textured.glsl"));
        m_texture = Texture2D::create(256, 256);
        
        uint32_t* pixel_data = new uint32_t[256 * 256];
        for (int y = 0; y < 256; y++)
        {
            for (int x = 0; x < 256; x++)
            {
                float u = (float)x / 256.0f;
                float v = (float)y / 256.0f;
                
                float wave = sin(u * 10.0f) * cos(v * 10.0f);
                wave = (wave + 1.0f) * 0.5f;
                
                uint8_t r = (uint8_t)(wave * 255);
                uint8_t g = (uint8_t)((1.0f - wave) * 255);
                uint8_t b = (uint8_t)((u + v) * 0.5f * 255);
                uint8_t a = 255;
                
                pixel_data[y * 256 + x] = (a << 24) | (b << 16) | (g << 8) | r;
            }
        }
        
        m_texture->set_data(pixel_data, 256 * 256 * 4);
        delete[] pixel_data;

        m_compute_shader = std::shared_ptr<Shader>(Shader::create("assets/texture_processor.glsl"));
        
        if (!m_compute_shader)
        {
            DONUT_WARN("Failed to create compute shader! Falling back to CPU texture generation.");
            m_use_compute_shader = false;
        }
        else
            DONUT_INFO("Compute shader created successfully!");
        
        m_processed_texture = Texture2D::create(256, 256);
    }
    
    void SimulationState::on_exit()
    {
    }
    
    void SimulationState::on_update(float delta_time)
    {
        m_delta_time = delta_time;
        handle_key_input(delta_time);
        update_texture();
    }
    
    void SimulationState::on_render()
    {
        Renderer::set_clear_color({ 0.0f, 0.0f, 0.0f, 1.0f });
        Renderer::clear();

        if (m_use_compute_shader && m_processed_texture)
        {
            m_processed_texture->bind(0);
            m_shader->set_int("u_Texture", 0);
        }
        else if (m_texture)
        {
            m_texture->bind(0);
            m_shader->set_int("u_Texture", 0);
        }

        glm::mat4 view_projection = m_camera ? m_camera->get_view_projection_matrix() : glm::mat4(1.0f);
        
        glm::vec3 center_pos(0.0f, 0.0f, 0.0f);
        glm::vec3 right_pos(2.0f, 0.0f, 0.0f);
        glm::vec3 left_pos(-2.0f, 0.0f, 0.0f);
        glm::vec3 up_pos(0.0f, 2.0f, 0.0f);
        glm::vec3 down_pos(0.0f, -2.0f, 0.0f);
        glm::vec3 forward_pos(0.0f, 0.0f, 2.0f);
        glm::vec3 back_pos(0.0f, 0.0f, -2.0f);
        
        glm::mat4 transform = glm::translate(glm::mat4(1.0f), center_pos);
        Renderer::submit(m_shader, m_vertex_array, transform, view_projection);
        transform = glm::translate(glm::mat4(1.0f), right_pos);
        Renderer::submit(m_shader, m_vertex_array, transform, view_projection);
        transform = glm::translate(glm::mat4(1.0f), left_pos);
        Renderer::submit(m_shader, m_vertex_array, transform, view_projection);
        transform = glm::translate(glm::mat4(1.0f), up_pos);
        Renderer::submit(m_shader, m_vertex_array, transform, view_projection);
        transform = glm::translate(glm::mat4(1.0f), down_pos);
        Renderer::submit(m_shader, m_vertex_array, transform, view_projection);
        transform = glm::translate(glm::mat4(1.0f), forward_pos);
        Renderer::submit(m_shader, m_vertex_array, transform, view_projection);
        transform = glm::translate(glm::mat4(1.0f), back_pos);
        Renderer::submit(m_shader, m_vertex_array, transform, view_projection);
    }
    
    void SimulationState::on_event(Event& event)
    {
        EventDispatcher dispatcher(event);
        
        dispatcher.dispatch<WindowResizeEvent>([this](WindowResizeEvent& e) 
        {
            if (m_camera)
                m_camera->set_projection(45.0f, (float)e.get_width() / (float)e.get_height(), 0.1f, 100.0f);
            return true;
        });
        
        dispatcher.dispatch<WindowFocusEvent>([this](WindowFocusEvent& e) 
        {
            m_first_mouse = true;
            m_mouse_dragging = false;
            return true;
        });
        
        dispatcher.dispatch<WindowLostFocusEvent>([this](WindowLostFocusEvent& e) 
        {
            m_mouse_dragging = false;
            return true;
        });
        
        dispatcher.dispatch<KeyPressedEvent>([this](KeyPressedEvent& e) 
        {
            m_keys[e.get_key_code()] = true;
            return true;
        });
        
        dispatcher.dispatch<KeyReleasedEvent>([this](KeyReleasedEvent& e) 
        {
            m_keys[e.get_key_code()] = false;
            return true;
        });
        
        dispatcher.dispatch<MouseMovedEvent>([this](MouseMovedEvent& e) 
        {
            if (!m_mouse_dragging)
                return true;
                
            if (m_first_mouse)
            {
                m_last_x = e.get_x();
                m_last_y = e.get_y();
                m_first_mouse = false;
            }

            float x_offset = m_last_x - e.get_x();
            float y_offset = e.get_y() - m_last_y;

            m_last_x = e.get_x();
            m_last_y = e.get_y();

            if (m_camera)
                m_camera->on_mouse_move(x_offset, y_offset);
            
            return true;
        });
        
        dispatcher.dispatch<MouseButtonPressedEvent>([this](MouseButtonPressedEvent& e) 
        {
            if (e.get_mouse_button() == GLFW_MOUSE_BUTTON_LEFT)
            {
                m_mouse_dragging = true;
                m_first_mouse = true;
            }
            return true;
        });
        
        dispatcher.dispatch<MouseButtonReleasedEvent>([this](MouseButtonReleasedEvent& e) 
        {
            if (e.get_mouse_button() == GLFW_MOUSE_BUTTON_LEFT)
                m_mouse_dragging = false;
            return true;
        });
    }
    
    void SimulationState::handle_key_input(float delta_time)
    {
        if (m_camera)
        {
            if (m_keys[GLFW_KEY_W])
                m_camera->move_forward(delta_time);
            if (m_keys[GLFW_KEY_S])
                m_camera->move_backward(delta_time);
            if (m_keys[GLFW_KEY_A])
                m_camera->move_left(delta_time);
            if (m_keys[GLFW_KEY_D])
                m_camera->move_right(delta_time);
            if (m_keys[GLFW_KEY_SPACE]) 
                m_camera->move_up(delta_time);
            if (m_keys[GLFW_KEY_LEFT_SHIFT])
                m_camera->move_down(delta_time);    
        }
        
        static bool c_key_pressed = false;
        if (m_keys[GLFW_KEY_C] && !c_key_pressed)
        {
            m_use_compute_shader = !m_use_compute_shader;
            DONUT_INFO("Compute shader: {}", (m_use_compute_shader ? "ON" : "OFF"));
            c_key_pressed = true;
        }
        if (!m_keys[GLFW_KEY_C])
            c_key_pressed = false;
        if (m_keys[GLFW_KEY_UP])
            m_compute_brightness += delta_time * 0.5f;
        if (m_keys[GLFW_KEY_DOWN])
            m_compute_brightness -= delta_time * 0.5f;
        if (m_keys[GLFW_KEY_LEFT])
            m_compute_contrast -= delta_time * 0.5f;
        if (m_keys[GLFW_KEY_RIGHT])
            m_compute_contrast += delta_time * 0.5f;
        if (m_keys[GLFW_KEY_Q])
            m_compute_saturation -= delta_time * 0.5f;
        if (m_keys[GLFW_KEY_E])
            m_compute_saturation += delta_time * 0.5f;
        
        m_compute_brightness = glm::clamp(m_compute_brightness, -1.0f, 1.0f);
        m_compute_contrast = glm::clamp(m_compute_contrast, 0.1f, 3.0f);
        m_compute_saturation = glm::clamp(m_compute_saturation, 0.0f, 2.0f);
        
        static float status_timer = 0.0f;
        status_timer += delta_time;
        if (status_timer > 2.0f)
        {
            DONUT_INFO("Compute Shader Status:\nBrightness: {}\nContrast: {}\nSaturation: {}",
                       m_compute_brightness, m_compute_contrast, m_compute_saturation);
            status_timer = 0.0f;
        }
    }
    
    void SimulationState::update_texture()
    {
        if (!m_texture)
            return;
            
        m_texture_time += m_delta_time;
        
        if (m_use_compute_shader && m_compute_shader && m_processed_texture)
        {
            m_compute_shader->bind();
            m_texture->bind(0);
            m_compute_shader->set_int("u_InputTexture", 0);
            m_processed_texture->bind_as_image(1, false);
            m_compute_shader->set_float2("u_TextureSize", glm::vec2(256.0f, 256.0f));
            m_compute_shader->set_float("u_Time", m_texture_time);
            m_compute_shader->set_float("u_Brightness", m_compute_brightness);
            m_compute_shader->set_float("u_Contrast", m_compute_contrast);
            m_compute_shader->set_float("u_Saturation", m_compute_saturation);
            m_compute_shader->dispatch(16, 16, 1);
            m_compute_shader->memory_barrier(IMAGE_ACCESS_BARRIER_BIT);
        }
        else
        {
            uint32_t* pixel_data = new uint32_t[256 * 256];
            for (int y = 0; y < 256; y++)
            {
                for (int x = 0; x < 256; x++)
                {
                    float time = m_texture_time * 2.0f;
                    float u = (float)x / 256.0f;
                    float v = (float)y / 256.0f;
                    
                    float wave = sin(u * 10.0f + time) * cos(v * 10.0f + time * 0.5f);
                    wave = (wave + 1.0f) * 0.5f;
                    
                    uint8_t r = (uint8_t)(wave * 255);
                    uint8_t g = (uint8_t)((1.0f - wave) * 255);
                    uint8_t b = (uint8_t)((u + v) * 0.5f * 255);
                    uint8_t a = 255;
                    
                    pixel_data[y * 256 + x] = (a << 24) | (b << 16) | (g << 8) | r;
                }
            }
            
            m_texture->set_data(pixel_data, 256 * 256 * 4);
            delete[] pixel_data;
        }
    }
};