// =============================================
//  Aster: light_manager.cpp
//  Copyright (c) 2020-2025 Anish Bhobe
// =============================================

#include "light_manager.h"

#include "aster/core/buffer.h"
#include "aster/systems/commit_manager.h"
#include "aster/systems/rendering_device.h"
#include "aster/systems/resource.h"
#include "glm/ext/matrix_transform.hpp"

// Static Checks

// Ensure layouts are exact.
static_assert(offsetof(DirectionalLight, m_Direction) == offsetof(Light, um_Direction));
static_assert(offsetof(DirectionalLight, m_Color_) == offsetof(Light, m_Color_));
static_assert(offsetof(DirectionalLight, m_Intensity) == offsetof(Light, m_Intensity));
static_assert(sizeof(DirectionalLight) <= sizeof(Light));

// Ensure layouts are exact.
static_assert(offsetof(PointLight, m_Position) == offsetof(Light, um_Position));
static_assert(offsetof(PointLight, m_Range) == offsetof(Light, m_Range));
static_assert(offsetof(PointLight, m_Color_) == offsetof(Light, m_Color_));
static_assert(offsetof(PointLight, m_Intensity) == offsetof(Light, m_Intensity));
static_assert(sizeof(PointLight) <= sizeof(Light));

// Ensure bitmask are in the right place.
static_assert((Light::TYPE_MASK & Light::TYPE_INVALID) == Light::TYPE_INVALID);
static_assert((Light::TYPE_MASK & Light::TYPE_DIRECTIONAL) == Light::TYPE_DIRECTIONAL);
static_assert((Light::TYPE_MASK & Light::TYPE_POINT) == Light::TYPE_POINT);
static_assert((Light::TYPE_MASK & Light::TYPE_SPOT) == Light::TYPE_SPOT);
static_assert(Light::COLOR_MASK == 0xFFFFFF00);

inline u32
ToColor32(vec4 const &col)
{
    u32 const r = static_cast<u32>(eastl::min(col.r, 1.0f) * 255.99f);
    u32 const g = static_cast<u32>(eastl::min(col.g, 1.0f) * 255.99f);
    u32 const b = static_cast<u32>(eastl::min(col.b, 1.0f) * 255.99f);
    u32 const a = static_cast<u32>(eastl::min(col.a, 1.0f) * 255.99f);

    return r << 24 | g << 16 | b << 8 | a;
}

inline u32
ToColor32(vec3 const &col)
{
    u32 const r = static_cast<u32>(eastl::min(col.r, 1.0f) * 255.99f);
    u32 const g = static_cast<u32>(eastl::min(col.g, 1.0f) * 255.99f);
    u32 const b = static_cast<u32>(eastl::min(col.b, 1.0f) * 255.99f);
    constexpr u32 a = 255;

    return r << 24 | g << 16 | b << 8 | a;
}

LightManager::LightManager(systems::RenderingDevice &device)
    : m_Device{&device}
    , m_DirectionalLightCount{}
    , m_PointLightCount{}
    , m_MetaInfo{}
    , m_GpuBufferCapacity_{0}
{
}

LightHandle
LightManager::AddDirectional(vec3 const &direction, vec3 const &color, f32 intensity)
{
    vec3 const normDirection = normalize(direction);
    if (m_DirectionalLightCount < m_MetaInfo.m_DirectionalLightMaxCount)
    {
        u16 index = 0;
        for (auto &light : m_Lights)
        {
            if (light.m_Range < 0)
            {
                u8 const gen = light.m_Color_ & Light::GEN_MASK;

                light.m_Color_ = (ToColor32(color) & Light::COLOR_MASK) | Light::TYPE_DIRECTIONAL | gen;
                light.m_Range = 1.0f;
                light.um_Direction = normDirection;
                light.m_Intensity = intensity;

                m_GpuBufferCapacity_ |= UPDATE_REQUIRED_BIT;

                return {Light::TYPE_DIRECTIONAL, gen, index};
            }
            ++index;
            assert(index < m_MetaInfo.m_DirectionalLightMaxCount); // Gap not found. But must exist
        }
    }

    // In case we will end up intersecting, we move point lights away. (2 at a time)
    if (m_DirectionalLightCount == m_MetaInfo.m_DirectionalLightMaxCount &&
        m_MetaInfo.m_DirectionalLightMaxCount == m_MetaInfo.m_PointLightOffset)
    {
        u16 const oldPointLightOffset = m_MetaInfo.m_PointLightOffset;
        u32 const pointLightMaxCount = m_MetaInfo.m_PointLightMaxCount;
        // Might cause a capacity increase, but I want to use that for my gpu buffer resize.
        m_Lights.push_back();
        m_Lights.push_back();

        if (m_MetaInfo.m_PointLightMaxCount > 0) // Edge Case: nullptr at size 0
        {
            Light *oldPointStart = m_Lights.data() + oldPointLightOffset;
            Light *oldPointEnd = oldPointStart + pointLightMaxCount;
            Light *newPointEnd = oldPointEnd + 2;

            static_assert(std::is_trivially_copyable_v<Light>);

            // Overlaps since 0 -> 2, 1 -> 3, 2 -> 4 (old 2 is lost)
            // Backward copy fixes this.
            std::copy_backward(oldPointStart, oldPointEnd, newPointEnd);
        }

        m_MetaInfo.m_PointLightOffset += 2;
    }

    constexpr u8 gen = 0; // New light
    m_Lights[m_DirectionalLightCount].m_Color_ = (ToColor32(color) & Light::COLOR_MASK) | Light::TYPE_DIRECTIONAL | gen;
    m_Lights[m_DirectionalLightCount].m_Range = 1.0f;
    m_Lights[m_DirectionalLightCount].um_Direction = normDirection;
    m_Lights[m_DirectionalLightCount].m_Intensity = intensity;
    u16 const index = m_DirectionalLightCount;

    ++m_DirectionalLightCount;
    ++m_MetaInfo.m_DirectionalLightMaxCount;

    return {Light::TYPE_DIRECTIONAL, gen, index};
}

LightHandle
LightManager::AddPoint(vec3 const &position, vec3 const &color, f32 const radius, f32 intensity)
{
    assert(m_PointLightCount <= m_MetaInfo.m_PointLightMaxCount);
    assert(radius >= 0.0f);
    if (m_PointLightCount < m_MetaInfo.m_PointLightMaxCount)
    {
        Light *light = m_Lights.data() + m_MetaInfo.m_PointLightOffset;
        for (u32 index = 0; index < m_MetaInfo.m_PointLightMaxCount; ++index)
        {
            if (light->m_Range < 0)
            {
                u8 const gen = light->m_Color_ & Light::GEN_MASK;

                light->m_Color_ = (ToColor32(color) & Light::COLOR_MASK) | Light::TYPE_POINT | gen;
                light->m_Range = radius;
                light->um_Position = position;
                light->m_Intensity = intensity;

                m_GpuBufferCapacity_ |= UPDATE_REQUIRED_BIT;

                return {Light::TYPE_POINT, gen, static_cast<u16>(index)};
            }
            ++light;
        }
        assert(false); // gap must exists.
        return {};
    }

    m_Lights.push_back();
    u16 const index = m_PointLightCount;

    Light *light = &m_Lights[index + m_MetaInfo.m_PointLightOffset];
    constexpr u8 gen = 0; // New light

    light->m_Color_ = (ToColor32(color) & Light::COLOR_MASK) | Light::TYPE_POINT | gen;
    light->m_Range = radius;
    light->um_Position = position;
    light->m_Intensity = intensity;

    ++m_PointLightCount;
    ++m_MetaInfo.m_PointLightMaxCount;

    m_GpuBufferCapacity_ |= UPDATE_REQUIRED_BIT;

    return {Light::TYPE_POINT, gen, index};
}

void
LightManager::Update()
{
    u16 const requiredBufferCapacity = eastl::min(static_cast<u16>(m_Lights.capacity()), MAX_LIGHTS);
    if ((m_GpuBufferCapacity_ & CAPACITY_MASK) < requiredBufferCapacity)
    {
        m_LightBuffer = m_Device->CreateStorageBuffer(requiredBufferCapacity * sizeof m_Lights[0], "Light Buffer");
        m_GpuBufferCapacity_ = requiredBufferCapacity | UPDATE_REQUIRED_BIT;

        m_MetaInfo.m_LightBuffer = m_LightBuffer->GetDeviceAddress();
    }
    if (m_GpuBufferCapacity_ & UPDATE_REQUIRED_BIT)
    {
        m_LightBuffer->Write(0, m_Lights.size() * sizeof m_Lights[0], m_Lights.data());
    }
}

void
LightManager::RemoveLight(LightHandle const handle)
{
    u8 const handleGen = handle.m_Generation;

    if (handle.m_Type == Light::TYPE_DIRECTIONAL)
    {
        Light *lightSlot = &m_Lights[handle.m_Index];
        u8 const slotGen = lightSlot->m_Color_ & Light::GEN_MASK;
        if (slotGen > handleGen)
        {
            WARN("Invalid handle gen: {} being freed. (slot gen: {})", handleGen, slotGen);
            return;
        }

        lightSlot->m_Range = -1.0f;
        lightSlot->m_Color_ = Light::TYPE_INVALID | (slotGen + 1) % Light::MAX_GEN;
        --m_DirectionalLightCount;
    }

    if (handle.m_Type == Light::TYPE_POINT)
    {
        Light *lightSlot = &m_Lights[handle.m_Index + m_MetaInfo.m_PointLightOffset];
        u8 const slotGen = lightSlot->m_Color_ & Light::GEN_MASK;
        if (slotGen > handleGen)
        {
            WARN("Invalid handle gen: {} being freed. (slot gen: {})", handleGen, slotGen);
            return;
        }

        lightSlot->m_Range = -1.0f;
        lightSlot->m_Color_ = Light::TYPE_INVALID | (slotGen + 1) % Light::MAX_GEN;
        --m_PointLightCount;
    }
}