#include <iostream>
#include <fstream>
#include <sstream>

#include <GLFW/glfw3.h>

#include "engine.h"
#include "rendering/vertex_buffer.h"
#include "rendering/index_buffer.h"

namespace Donut
{
    Engine::Engine()
        : m_sag_a(glm::vec3(0.0f, 0.0f, 0.0f), 8.54e36f)
    {
        GLFWwindow* window = static_cast<GLFWwindow*>(Application::get().get_window().get_native_window());
        glfwGetFramebufferSize(window, &m_width, &m_height);
        
        m_camera.set_camera_mode(CameraMode::Orbital);
        m_camera.set_orbital_radius(6.34194e10);
        m_camera.set_orbital_limits(1e10, 1e12);
        m_camera.set_orbital_speed(0.01f);
        m_camera.set_zoom_speed(25e9f);
        
        m_objects = 
        {
            { glm::vec4(4e11f, 0.00f, 0.00f, 4e10f),       glm::vec4(1, 1, 0, 1), static_cast<float>(1.98892e30)    },
            { glm::vec4(0.00f, 0.00f, 4e11f, 4e10f),       glm::vec4(1, 0, 0, 1), static_cast<float>(1.98892e30)    },
            { glm::vec4(0.00f, 0.00f, 0.00f, m_sag_a.m_rs), glm::vec4(0, 0, 0, 1), static_cast<float>(m_sag_a.m_mass) }
        };

        m_shader_program  = Ref<Shader>(Shader::create("assets/shaders/compute_texture_quad.glsl"));
        m_compute_program = create_compute_program("assets/shaders/geodesic.glsl");
        
        m_camera_ubo = UniformBuffer::create(128, 1);
        m_disk_ubo   = UniformBuffer::create(sizeof(float) * 4, 2);
        
        uint32_t obj_ubo_size = sizeof(int) + 3 * sizeof(float)
            + 16 * (sizeof(glm::vec4) + sizeof(glm::vec4))
            + 16 * sizeof(float);
        m_objects_ubo = UniformBuffer::create(obj_ubo_size, 3);

        auto result = quad_vao();
        m_quad_vao = result.first;
        m_texture = result.second;
    }

    void Engine::update_window_dimensions()
    {
        GLFWwindow* window = static_cast<GLFWwindow*>(Application::get().get_window().get_native_window());
        int old_width = m_width;
        int old_height = m_height;
        int old_compute_height = m_compute_height;
        
        glfwGetFramebufferSize(window, &m_width, &m_height);
        
        if (old_width         != m_width  || 
            old_height        != m_height || 
            old_compute_height != m_compute_height)
            update_compute_dimensions();
    }

    void Engine::update_performance(float delta_time)
    {
        if (delta_time > 0.0f)
            m_current_fps = 1.0f / delta_time;
    }

    void Engine::update_compute_dimensions()
    {
        m_texture = Texture2D::create(get_compute_width(), m_compute_height);
    }

    void Engine::draw_full_screen_quad()
    {
        RenderCommand::set_viewport(0, 0, m_width, m_height);
        
        m_shader_program->bind();
        m_quad_vao->bind();

        m_texture->bind(0);
        m_shader_program->set_int("u_ScreenTexture", 0);

        RenderCommand::disable_depth_test();
        RenderCommand::draw_arrays(6);
        RenderCommand::enable_depth_test();
    }

    void Engine::dispatch_compute(const Camera& cam)
    {
        int cw = get_compute_width();
        int ch = m_compute_height;

        m_texture->set_data(nullptr, cw * ch * 4);

        m_compute_program->bind();
        upload_camera_ubo(cam);
        upload_disk_ubo();
        upload_objects_ubo(m_objects);
        m_texture->bind_as_image(0, false);

        uint32_t groups_x = static_cast<uint32_t>(std::ceil(cw / 16.0f));
        uint32_t groups_y = static_cast<uint32_t>(std::ceil(ch / 16.0f));
        m_compute_program->dispatch(groups_x, groups_y, 1);
        m_compute_program->memory_barrier(IMAGE_ACCESS_BARRIER_BIT);
    }

    void Engine::upload_camera_ubo(const Camera& cam)
    {
        struct UBOData
        {
            glm::vec3 pos;     float _pad0;
            glm::vec3 right;   float _pad1;
            glm::vec3 up;      float _pad2;
            glm::vec3 forward; float _pad3;
            float tan_half_fov;
            float aspect;
            bool  moving;
            int   _pad4;
        } data;

        glm::vec3 fwd   = glm::normalize(cam.get_orbital_target() - cam.get_orbital_position());
        glm::vec3 up    = glm::vec3(0, 1, 0);
        glm::vec3 right = glm::normalize(glm::cross(fwd, up));
        up = glm::cross(right, fwd);

        data.pos        = cam.get_orbital_position();
        data.right      = right;
        data.up         = up;
        data.forward    = fwd;
        data.tan_half_fov = static_cast<float>(tan(glm::radians(60.0f * 0.5f)));
        data.aspect     = static_cast<float>(get_compute_width()) / static_cast<float>(m_compute_height);
        data.moving     = cam.is_dragging() || cam.is_panning();

        m_camera_ubo->set_data(&data, sizeof(UBOData));
        m_camera_ubo->bind(1);
    }

    void Engine::upload_objects_ubo(const std::vector<ObjectData>& objs)
    {
        struct UBOData
        {
            int num_objects;
            float _pad0, _pad1, _pad2;
            glm::vec4 pos_radius[16];
            glm::vec4 color[16];
            float mass[16];
        } data;

        size_t count = std::min(objs.size(), size_t(16));
        data.num_objects = static_cast<int>(count);

        for (size_t i = 0; i < count; ++i)
        {
            data.pos_radius[i] = objs[i].m_pos_radius;
            data.color[i] = objs[i].m_color;
            data.mass[i] = objs[i].m_mass;
        }

        m_objects_ubo->set_data(&data, sizeof(data));
        m_objects_ubo->bind(3);
    }

    void Engine::upload_disk_ubo()
    {
        float r1 = static_cast<float>(m_sag_a.m_rs * 2.2);
        float r2 = static_cast<float>(m_sag_a.m_rs * 5.2);
        float num = 2.0f;
        float thickness = 1e9f;
        float disk_data[4] = { r1, r2, num, thickness };

        m_disk_ubo->set_data(disk_data, sizeof(disk_data));
        m_disk_ubo->bind(2);
    }

    void Engine::update_physics(float delta_time)
    {
        for (auto& obj : m_objects)
        {
            for (auto& obj2 : m_objects)
            {
                if (&obj == &obj2) continue;
                float dx = obj2.m_pos_radius.x - obj.m_pos_radius.x;
                float dy = obj2.m_pos_radius.y - obj.m_pos_radius.y;
                float dz = obj2.m_pos_radius.z - obj.m_pos_radius.z;
                float distance = sqrt(dx * dx + dy * dy + dz * dz);
                if (distance > 0)
                {
                    std::vector<double> direction = {dx / distance, dy / distance, dz / distance};
                    double Gforce = (G * obj.m_mass * obj2.m_mass) / (distance * distance);
                    double acc1 = Gforce / obj.m_mass;
                    std::vector<double> acc = {direction[0] * acc1, direction[1] * acc1, direction[2] * acc1};
                    
                    if (m_gravity)
                    {
                        obj.m_velocity.x += static_cast<float>(acc[0]);
                        obj.m_velocity.y += static_cast<float>(acc[1]);
                        obj.m_velocity.z += static_cast<float>(acc[2]);

                        obj.m_pos_radius.x += static_cast<float>(obj.m_velocity.x);
                        obj.m_pos_radius.y += static_cast<float>(obj.m_velocity.y);
                        obj.m_pos_radius.z += static_cast<float>(obj.m_velocity.z);
                    }
                }
            }
        }
    }

    void Engine::render_scene()
    {
        RenderCommand::clear();
        m_shader_program->bind();
        m_quad_vao->bind();
        m_texture->bind(0);
        RenderCommand::draw_arrays(6);
    }

    Ref<Shader> Engine::create_compute_program(const char* path)
    {
        std::ifstream in(path);
        if(!in.is_open())
        {
            std::cerr << "Failed to open compute shader: " << path << "\n";
            return nullptr;
        }
        std::stringstream ss;
        ss << in.rdbuf();
        std::string src_str = ss.str();
        return Ref<Shader>(Shader::create_compute("ComputeShader", src_str));
    }

    std::pair<Ref<VertexArray>, Ref<Texture2D>> Engine::quad_vao()
    {
        float quad_vertices[] = 
        {
            -1.0f,  1.0f,  0.0f, 1.0f,
            -1.0f, -1.0f,  0.0f, 0.0f,
             1.0f, -1.0f,  1.0f, 0.0f,
            -1.0f,  1.0f,  0.0f, 1.0f,
             1.0f, -1.0f,  1.0f, 0.0f,
             1.0f,  1.0f,  1.0f, 1.0f
        };

        auto vertex_buffer = Ref<VertexBuffer>(VertexBuffer::create(quad_vertices, static_cast<uint32_t>(sizeof(quad_vertices))));
        VertexBufferLayout layout;
        layout.push<float>(2);
        layout.push<float>(2);
        vertex_buffer->set_layout(layout);

        auto vertex_array = Ref<VertexArray>(VertexArray::create());
        vertex_array->add_vertex_buffer(vertex_buffer);
        auto texture = Texture2D::create(get_compute_width(), m_compute_height);
        
        return { vertex_array, texture };
    }
}