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

#include "light_manager.h"

#include "buffer.h"
#include "glm/ext/matrix_transform.hpp"

struct Light
{
    union {
        vec3 um_Position;
        vec3 um_Direction;
    };
    f32 m_Range;  // < 0.0 for invalid
    u32 m_Color_; // LSB is used for flags. (R G B Flags)
    f32 m_Intensity;

    constexpr static u32 MAX_GEN = 0x40;
    constexpr static u32 GEN_MASK = MAX_GEN - 1;

    constexpr static u32 TYPE_MASK = 0xC0;
    constexpr static u32 TYPE_INVALID = 0x0;
    constexpr static u32 TYPE_DIRECTIONAL = 1 << 6;
    constexpr static u32 TYPE_POINT = 2 << 6;
    constexpr static u32 TYPE_SPOT = 3 << 6; // Currently Unused

    constexpr static u32 COLOR_MASK = ~(GEN_MASK | TYPE_MASK);
};

// 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(const vec4 &col)
{
    const u32 r = Cast<u32>(eastl::min(col.r, 1.0f) * 255.99f);
    const u32 g = Cast<u32>(eastl::min(col.g, 1.0f) * 255.99f);
    const u32 b = Cast<u32>(eastl::min(col.b, 1.0f) * 255.99f);
    const u32 a = Cast<u32>(eastl::min(col.a, 1.0f) * 255.99f);

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

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

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

LightManager::LightManager(GpuResourceManager *resourceManager)
    : m_ResourceManager{resourceManager}
    , m_DirectionalLightCount{}
    , m_PointLightCount{}
    , m_MetaInfo{}
    , m_GpuBufferCapacity_{0}
{
}

LightManager::~LightManager()
{
    m_ResourceManager->Release(m_MetaInfo.m_LightBuffer);
}

LightManager::LightManager(LightManager &&other) noexcept
    : m_ResourceManager(other.m_ResourceManager)
    , m_Lights(std::move(other.m_Lights))
    , m_DirectionalLightCount(other.m_DirectionalLightCount)
    , m_PointLightCount(other.m_PointLightCount)
    , m_MetaInfo(other.m_MetaInfo)
    , m_GpuBufferCapacity_(other.m_GpuBufferCapacity_)
{
    other.m_MetaInfo.m_LightBuffer = {};
}

LightManager &
LightManager::operator=(LightManager &&other) noexcept
{
    if (this == &other)
        return *this;
    m_ResourceManager = other.m_ResourceManager;
    m_Lights = std::move(other.m_Lights);
    m_DirectionalLightCount = other.m_DirectionalLightCount;
    m_PointLightCount = other.m_PointLightCount;
    m_MetaInfo = other.m_MetaInfo;
    other.m_MetaInfo.m_LightBuffer = {};
    m_GpuBufferCapacity_ = other.m_GpuBufferCapacity_;
    return *this;
}

LightHandle
LightManager::AddDirectional(const vec3 &direction, const vec3 &color, f32 intensity)
{
    const vec3 normDirection = normalize(direction);
    if (m_DirectionalLightCount < m_MetaInfo.m_DirectionalLightMaxCount)
    {
        u16 index = 0;
        for (auto &light : m_Lights)
        {
            if (light.m_Range < 0)
            {
                const u8 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)
    {
        const u16 oldPointLightOffset = m_MetaInfo.m_PointLightOffset;
        const u32 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 *newPointStart = oldPointStart + 2;

            static_assert(std::is_trivially_copyable_v<Light>);
            memcpy(newPointStart, oldPointStart, pointLightMaxCount * sizeof *newPointStart);
        }

        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;
    const u16 index = m_DirectionalLightCount;

    ++m_DirectionalLightCount;
    ++m_MetaInfo.m_DirectionalLightMaxCount;

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

LightHandle
LightManager::AddPoint(const vec3 &position, const vec3 &color, const f32 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)
            {
                const u8 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, Cast<u16>(index)};
            }
            ++light;
        }
        assert(false); // gap must exists.
        return {};
    }

    m_Lights.push_back();
    const u16 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()
{
    const u16 requiredBufferCapacity = eastl::min(Cast<u16>(m_Lights.capacity()), MAX_LIGHTS);
    if ((m_GpuBufferCapacity_ & CAPACITY_MASK) < requiredBufferCapacity)
    {
        StorageBuffer newBuffer;
        newBuffer.Init(m_ResourceManager->m_Device, requiredBufferCapacity * sizeof m_Lights[0], true, "Light Buffer");
        m_GpuBufferCapacity_ = requiredBufferCapacity | UPDATE_REQUIRED_BIT;

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

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

    if (handle.m_Type == Light::TYPE_DIRECTIONAL)
    {
        Light *lightSlot = &m_Lights[handle.m_Index];
        const u8 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];
        const u8 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;
    }
}