#include "HydraDirectionalShadowSourceSystem.h"
#include "../../engine/HydraEngine.h"
#include "../../camera/HydraCamera.h"
#include "../HydraECS.h"
#include "../components/HydraDirectionalShadowSourceComponent.h"
#include "../components/HydraPositionComponent.h"
#include "../components/HydraLightComponent.h"

#include "../../buffer/HydraShadowBuffer.h"
#include "../../buffer/HydraTextureBufferManager.h"

#include "../../scheduler/HydraTaskScheduler.h"
#include "../../task/buffer/HydraCreateGPUBufferTask.hpp"
#include "../../task/buffer/HydraUpdateWholeGPUBufferTask.hpp"
#include "../../task/buffer/HydraUpdatePartialGPUBufferTask.hpp"
#include "../../task/buffer/HydraCreateUnBoundTextureTask.hpp"
#include "../../task/renderer/HydraCreateCascadeShadowFrameBuffer.hpp"
#include "../../task/renderer/HydraCreateShadowFrameBuffer.hpp"
#include "../../actor/HydraActorManager.h"

void HydraDirectionalShadowSourceSystem::Initialise()
{
    HydraTaskScheduler *task_scheduler = GetEngine()->TaskScheduler();

    for(uint32_t i = 0; i < MAX_SHADOWS; i++)
    {
        m_available_shadows.push(i);
    }

    m_shadows.resize(MAX_SHADOWS);

    HydraID create_gpu_buffer_task_id = task_scheduler->CreateTask<HydraCreateGPUBufferTask>(HydraThreadAffinity::Render);
    HydraCreateGPUBufferTask *task = static_cast<HydraCreateGPUBufferTask *>(task_scheduler->GetTask(create_gpu_buffer_task_id));
    task->Intialise("ShadowBuffer", 1, MAX_SHADOWS * sizeof(HydraDirectionalShadowSourceComponent), &m_shadow_buffer_id, HYDRA_BUFFER_TYPE::HYDRA_SSBO_BUFFER);
    task_scheduler->WaitForTask(create_gpu_buffer_task_id);
}

void HydraDirectionalShadowSourceSystem::InitialiseComponent(HydraID entity_id)
{
    HydraEngine *engine = GetEngine();
    HydraECS *ecs = engine->ECS();
    HydraLightComponent &light_comp = ecs->GetComponent<HydraLightComponent>(entity_id);

    {
        std::lock_guard<std::mutex> lock(m_mutex);
        assert(m_available_shadows.size() > 0 && "Run out of shadows");
        uint32_t shad_id = m_available_shadows.front();
        m_available_shadows.pop();
        light_comp.shadow_map_id = shad_id;
        light_comp.shadow_caster = 1;
    }
    

    _InitialiseShadowMap(entity_id);

}

void HydraDirectionalShadowSourceSystem::UpdateSystem(float delta_t_seconds)
{
    HydraEngine *engine = GetEngine();
    HydraECS *ecs = engine->ECS();
    std::vector<HydraDirectionalShadowSourceComponent> shadows;
    for (HydraID entity_id : Entities)
    {

        HydraLightComponent &light_comp = ecs->GetComponent<HydraLightComponent>(entity_id);

        _UpdateShadowMap(entity_id);

        HydraDirectionalShadowSourceComponent shad_comp = ecs->GetComponent<HydraDirectionalShadowSourceComponent>(entity_id);
        m_shadows[light_comp.shadow_map_id] = shad_comp;
    }
    _WriteToGPU();
}

void HydraDirectionalShadowSourceSystem::_UpdateShadowMap(HydraID entity_id)
{
    HydraEngine *engine = GetEngine();
    HydraECS *ecs = engine->ECS();
    HydraPositionComponent pos_comp = ecs->GetComponent<HydraPositionComponent>(entity_id);
    HydraDirectionalShadowSourceComponent &shad_comp = ecs->GetComponent<HydraDirectionalShadowSourceComponent>(entity_id);

    float view_size = 4.0f;
    float view_near = 0.1f;
    float view_far = 4.0f;
    
    glm::mat4 light_rot = glm::eulerAngleXYZ(pos_comp.orientation.x, pos_comp.orientation.y, pos_comp.orientation.z);
    glm::vec3 light_dir = glm::normalize(glm::vec3(light_rot * glm::vec4(0.0f, 0.0f, -1.0f, 0.0f)));

    glm::vec3 camera_pos = GetEngine()->Camera()->GetPosition();

    for (uint32_t i = 0; i < 6; i++)
    {
        glm::vec3 light_pos = camera_pos;//- (light_dir * 2.0f);          
        glm::vec3 target_pos = light_pos + light_dir;
        glm::vec3 up_vector = glm::vec3(0.0f, 1.0f, 0.0f);

        if (glm::abs(glm::dot(light_dir, up_vector)) > 0.99f) 
        {
            up_vector = glm::vec3(1.0f, 0.0f, 0.0f); // Prevent gimbal lock
        }

        glm::mat4 view = glm::lookAt(light_pos, target_pos, up_vector);         
        glm::mat4 proj = glm::ortho(-view_size, view_size, -view_size, view_size, view_near, view_far);          

        // Save data to component
        shad_comp.far_plane[i] = view_far;         
        shad_comp.near_plane[i] = view_near;         
        shad_comp.light_space_proj[i] = proj;         
        shad_comp.light_space_view[i] = view;          

        // Scale bounds exponentially for the next cascade tier
        view_size *= 3.0f;     
        view_far  *= 3.0f;        
    }
}


void HydraDirectionalShadowSourceSystem::_InitialiseShadowMap(HydraID entity_id)
{
    HydraEngine *engine = GetEngine();
    HydraECS *ecs = engine->ECS();
    HydraTextureBufferManager *tex_man = engine->TextureBufferManager();
    HydraDirectionalShadowSourceComponent &shad_comp = ecs->GetComponent<HydraDirectionalShadowSourceComponent>(entity_id);
    HydraLightComponent &light_comp = ecs->GetComponent<HydraLightComponent>(entity_id);

    shad_comp.entity_id = entity_id;
    shad_comp.cascade_count = 6;
    
    HydraTaskScheduler *task_scheduler = GetEngine()->TaskScheduler();

    uint32_t depth_texture_id = 0;
    HydraID create_framebuffer_task_id = task_scheduler->CreateTask<HydraCreateCascadeShadowFrameBuffer>(HydraThreadAffinity::Render);
    HydraCreateCascadeShadowFrameBuffer *fb_task = static_cast<HydraCreateCascadeShadowFrameBuffer *>(task_scheduler->GetTask(create_framebuffer_task_id));
    fb_task->Initialise(&shad_comp.shadow_texture_id, &depth_texture_id, &shad_comp.frame_buffer_id);
    
    task_scheduler->WaitForTask(create_framebuffer_task_id);
    shad_comp.shadow_texture_id = depth_texture_id;
}

void HydraDirectionalShadowSourceSystem::_WriteToGPU()
{
    HydraEngine *engine = GetEngine();
    HydraECS *ecs = engine->ECS();

    void *buffer = static_cast<void *>(m_shadows.data());
    HydraTaskScheduler *task_scheduler = GetEngine()->TaskScheduler();

    HydraID update_buffer_task_id = task_scheduler->CreateTask<HydraUpdateWholeGPUBufferTask>(HydraThreadAffinity::Render);
    HydraUpdateWholeGPUBufferTask *task = static_cast<HydraUpdateWholeGPUBufferTask *>(task_scheduler->GetTask(update_buffer_task_id));
    task->Intialise(m_shadow_buffer_id,  buffer);
    task_scheduler->WaitForTask(update_buffer_task_id);
}