#include <Core/HAL/Display/Win32/Win32Window.h>
#include <Core/HAL/Display/Window.h>
#include <Core/ImGui/ImGuiService.h>
#include <Core/ImGui/ImGuiTests.h>
#include <Core/Logging/LogManager.h>
#include <Core/Memory/MemoryArena.h>

#ifdef JULIET_ENABLE_IMGUI

#include <backends/imgui_impl_win32.h>
#include <imgui.h>

extern IMGUI_IMPL_API LRESULT ImGui_ImplWin32_WndProcHandler(HWND hWnd, UINT msg, WPARAM wParam, LPARAM lParam);

namespace Juliet::ImGuiService
{
    namespace
    {
        ImGuiContext* g_ImGuiContext = nullptr;
        bool          g_Initialized  = false;

        struct FreeNode
        {
            size_t    Size;
            FreeNode* Next;
        }* FirstFreeNode = nullptr;

        // Dedicated Paged Arena for ImGui
        // Sharing the same underlying Engine Pool for blocks, but separate Arena chain.
        Arena* g_ImGuiArena = {};

        void* ImGuiAllocWrapper(size_t size, void* /*user_data*/)
        {
            Assert(size); // I trust Dear Imgui but just in case

            // Let's make FreeNode a part of each alloc.
            size_t     sizeOfHeader = sizeof(FreeNode);
            size_t     totalSize    = AlignPow2<size_t>(size + sizeOfHeader, 16);
            FreeNode** previous     = &FirstFreeNode;
            FreeNode*  current      = FirstFreeNode;

            // Find a free node if there is one big enough
            while (current != nullptr)
            {
                if (current->Size >= totalSize)
                {
                    *previous = current->Next;
                    return current + 1;
                }
                previous = &current->Next;
                current  = current->Next;
            }

            auto* ptr = ArenaPushSize(g_ImGuiArena, totalSize, 8, false JULIET_DEBUG_PARAM("ImGuiAlloc {}", totalSize));

            FreeNode* node = static_cast<FreeNode*>(ptr);
            node->Size     = totalSize;
            node->Next     = nullptr;

            return node + 1;
        }

        void ImGuiFreeWrapper(void* ptr, void* /*user_data*/)
        {
            Assert(ptr); // I trust Dear Imgui but just in case

            FreeNode* node = reinterpret_cast<FreeNode*>(static_cast<char*>(ptr) - sizeof(FreeNode));
            node->Next     = FirstFreeNode;
            FirstFreeNode  = node;
        }
    } // namespace

    void Initialize(NonNullPtr<Window> window)
    {
        Assert(!g_Initialized);

        // Initialize ImGui Arena using Engine Pool
        g_ImGuiArena = ArenaAllocate({ .Name = "Juliet" });

        // Setup Allocator
        ImGui::SetAllocatorFunctions(ImGuiAllocWrapper, ImGuiFreeWrapper, nullptr);

        IMGUI_CHECKVERSION();
        g_ImGuiContext = ImGui::CreateContext();
        ImGuiIO& io    = ImGui::GetIO();
        io.ConfigFlags |= ImGuiConfigFlags_NavEnableKeyboard; // Enable Keyboard Controls
        // io.ConfigFlags |= ImGuiConfigFlags_NavEnableGamepad;      // Enable Gamepad Controls

        ImGui::StyleColorsDark();

        // Platform Init
        auto* win32State = static_cast<Win32::Window32State*>(window->State);
        ImGui_ImplWin32_Init(win32State->Handle);

        g_Initialized = true;
    }

    void Shutdown()
    {
        Assert(g_Initialized);

        ImGui_ImplWin32_Shutdown();
        ImGui::DestroyContext(g_ImGuiContext);
        g_ImGuiContext = nullptr;
        g_Initialized  = false;
    }

    void NewFrame()
    {
        Assert(g_Initialized);

        ImGui_ImplWin32_NewFrame();
        ImGui::NewFrame();
    }

    void Render()
    {
        Assert(g_Initialized);

        ImGui::Render();
    }

    bool IsInitialized()
    {
        return g_Initialized;
    }

    ImGuiContext* GetContext()
    {
        return g_ImGuiContext;
    }

    void RunTests()
    {
        printf("ImGuiService: Running Unit Tests...\n");
        Juliet::UnitTest::TestImGui();
    }
} // namespace Juliet::ImGuiService

#endif // JULIET_ENABLE_IMGUI