#include <Core/Common/String.h>
#include <Core/Container/Vector.h>
#include <Core/Memory/MemoryArena.h> // Added for Arena definition
#include <Core/Memory/MemoryArenaDebug.h>
#include <Engine/Debug/MemoryDebugger.h>

#include <imgui.h>

#if JULIET_DEBUG

namespace Juliet
{
    Arena* GetGlobalArenaListHead();
} // namespace Juliet

namespace Juliet::Debug
{
    namespace
    {
        struct ArenaDebugState
        {
            float            Zoom                   = 1.0f;
            ArenaAllocation* SelectedAlloc          = nullptr;
            bool             ScrollVisualToSelected = false;
            bool             ScrollListToSelected   = false;
        };

        struct PagedArenaDebugState
        {
            float           Zoom                   = 1.0f;
            ArenaDebugInfo* SelectedAlloc          = nullptr;
            bool            ScrollVisualToSelected = false;
            bool            ScrollListToSelected   = false;
        };

        struct PagedArenaStateEntry
        {
            const Arena*         Arena;
            PagedArenaDebugState State;
        };

        // Static Vector for lookup
        VectorArena<PagedArenaStateEntry, Kilobytes(1024)> s_PagedArenaStates;

        PagedArenaDebugState& GetPagedArenaState(const Arena* arena)
        {
            if (s_PagedArenaStates.Arena == nullptr)
            {
                Arena* pagedArenaStates = ArenaAllocate({} JULIET_DEBUG_ONLY(, "DebugState States"));
                s_PagedArenaStates.Create(pagedArenaStates);
            }

            for (auto& entry : s_PagedArenaStates)
            {
                if (entry.Arena == arena)
                {
                    return entry.State;
                }
            }

            PagedArenaStateEntry newEntry;
            newEntry.Arena = arena;
            // Default construct state
            newEntry.State = PagedArenaDebugState();

            s_PagedArenaStates.PushBack(newEntry);
            return s_PagedArenaStates.Last()->State;
        }

        // Generate a unique color per allocation, using tag + offset + size
        // so entries with same tag name get distinct colors
        uint32 GetColorForAllocation(const char* tag, size_t offset, size_t size)
        {
            uint32 hash = 2166136261u;
            if (tag)
            {
                for (const char* s = tag; *s; ++s)
                {
                    hash ^= static_cast<uint8>(*s);
                    hash *= 16777619u;
                }
            }
            // Mix in offset and size to differentiate same-tag entries
            hash ^= static_cast<uint32>(offset);
            hash *= 16777619u;
            hash ^= static_cast<uint32>(size);
            hash *= 16777619u;

            float h = static_cast<float>(hash % 360) / 360.0f;
            float s = 0.5f + static_cast<float>((hash >> 12) % 30) / 100.0f;
            float v = 0.65f + static_cast<float>((hash >> 20) % 25) / 100.0f;
            return ImColor::HSV(h, s, v);
        }

        void DrawPagedArena(const Arena* arena)
        {
            ImGui::PushID(arena);
            const char*           name  = (arena->Name && arena->Name[0] != '\0') ? arena->Name : "Unnamed Arena";
            PagedArenaDebugState& state = GetPagedArenaState(arena);

            if (ImGui::CollapsingHeader(name))
            {
                // Collect Blocks (Pages) in order (Oldest -> Newest)
                // Arena->Previous chain goes Newest -> Oldest.
                static VectorArena<const Arena*> blocks;
                if (blocks.Arena == nullptr)
                {
                    Arena* pagedArenaStates = ArenaAllocate({} JULIET_DEBUG_ONLY(, "DebugState Blocks"));
                    blocks.Create(pagedArenaStates);
                }
                blocks.Clear();

                for (const Arena* curr = arena->Current; curr != nullptr; curr = curr->Previous)
                {
                    blocks.PushBack(curr);
                }

                // Reverse manually
                size_t count = blocks.Count;
                for (size_t i = 0; i < count / 2; ++i)
                {
                    const Arena* temp     = blocks[i];
                    blocks[i]             = blocks[count - 1 - i];
                    blocks[count - 1 - i] = temp;
                }

                // Calculate Stats
                size_t totalCapacity = 0;
                size_t totalUsed     = 0;

                for (const Arena* blk : blocks)
                {
                    totalCapacity += blk->Reserved;
                    totalUsed += blk->Position;
                }

                ImGui::Text("Used: %zu / %zu bytes (%zu pages)", totalUsed, totalCapacity, blocks.Size());
                ImGui::SliderFloat("Zoom", &state.Zoom, 0.1f, 1000000.0f, "%.2f", ImGuiSliderFlags_Logarithmic);

                // --- Visual View ---
                ImGui::Separator();
                ImGui::Text("Visual Map");

                float blockHeight       = 24.0f;
                float blockSpacing      = 4.0f;
                float requiredVisHeight = static_cast<float>(blocks.Size()) * (blockHeight + blockSpacing) + 30.0f;
                // constrains
                if (requiredVisHeight > 300.0f)
                {
                    requiredVisHeight = 300.0f;
                }
                if (requiredVisHeight < 50.0f)
                {
                    requiredVisHeight = 50.0f;
                }

                if (ImGui::BeginChild("VisualMap", ImVec2(0, requiredVisHeight), true,
                                      ImGuiWindowFlags_HorizontalScrollbar | ImGuiWindowFlags_NoMove | ImGuiWindowFlags_NoScrollbar))
                {
                    ImGui::SetScrollY(0.0f);

                    ImDrawList* dl         = ImGui::GetWindowDrawList();
                    float       availWidth = ImGui::GetContentRegionAvail().x;
                    ImVec2      startPos   = ImGui::GetCursorScreenPos();

                    float virtualWidth = availWidth * state.Zoom;
                    if (virtualWidth < availWidth)
                    {
                        virtualWidth = availWidth;
                    }
                    ImGui::Dummy(ImVec2(virtualWidth, requiredVisHeight + 1.0f));

                    bool isMapHovered = ImGui::IsWindowHovered();

                    // Zoom with pivot
                    if (isMapHovered)
                    {
                        float wheel = ImGui::GetIO().MouseWheel;
                        if (wheel != 0.0f)
                        {
                            float mouseXScreen     = ImGui::GetMousePos().x;
                            float activeScrollX    = ImGui::GetScrollX();
                            float viewportLeft     = ImGui::GetWindowPos().x;
                            float mouseRelViewport = mouseXScreen - viewportLeft;

                            float virtualWidthOld = std::max(availWidth * state.Zoom, availWidth);
                            float mouseContentPos = activeScrollX + mouseRelViewport;
                            float ratio           = mouseContentPos / virtualWidthOld;

                            state.Zoom += wheel * 0.2f * state.Zoom;
                            state.Zoom = std::max(state.Zoom, 0.1f);
                            state.Zoom = std::min(state.Zoom, 1000000.0f);

                            float virtualWidthNew = std::max(availWidth * state.Zoom, availWidth);
                            float desiredScrollX  = (ratio * virtualWidthNew) - mouseRelViewport;
                            ImGui::SetScrollX(desiredScrollX);
                        }
                    }

                    // Pan
                    if (isMapHovered && (ImGui::IsMouseDragging(ImGuiMouseButton_Left) || ImGui::IsMouseDragging(ImGuiMouseButton_Right) ||
                                         ImGui::IsMouseDragging(ImGuiMouseButton_Middle)))
                    {
                        ImGui::SetScrollX(ImGui::GetScrollX() - ImGui::GetIO().MouseDelta.x);
                    }

                    // Jump to Selected (Sync from List)
                    if (state.ScrollVisualToSelected && state.SelectedAlloc)
                    {
                        for (const Arena* searchBlk : blocks)
                        {
                            ArenaDebugInfo* search = searchBlk->FirstDebugInfo;
                            bool            found  = false;
                            while (search)
                            {
                                if (search == state.SelectedAlloc)
                                {
                                    found = true;
                                    break;
                                }
                                search = search->Next;
                            }

                            if (found)
                            {
                                float virtualWidthLocal = std::max(availWidth * state.Zoom, availWidth);
                                if (searchBlk->Reserved > 0)
                                {
                                    double scale =
                                        static_cast<double>(virtualWidthLocal) / static_cast<double>(searchBlk->Reserved);
                                    float xStart = static_cast<float>(static_cast<double>(state.SelectedAlloc->Offset) * scale);
                                    float centerOffset = availWidth * 0.5f;
                                    ImGui::SetScrollX(xStart - centerOffset);
                                }
                                state.ScrollVisualToSelected = false;
                                break;
                            }
                        }
                    }

                    ImVec2 pos = startPos;
                    for (const Arena* blk : blocks)
                    {
                        ImGui::PushID(blk);
                        ImVec2 rectMin = pos;
                        ImVec2 rectMax = ImVec2(pos.x + virtualWidth, pos.y + blockHeight);

                        dl->AddRect(rectMin, rectMax, IM_COL32(100, 100, 100, 255));

                        size_t dataSize = blk->Reserved - k_ArenaHeaderSize;
                        if (dataSize > 0)
                        {
                            double scale = static_cast<double>(virtualWidth) / static_cast<double>(blk->Reserved);

                            // Draw Arena Header
                            {
                                float hdrWidth = static_cast<float>(static_cast<double>(k_ArenaHeaderSize) * scale);
                                hdrWidth       = std::max(hdrWidth, 1.0f);
                                ImVec2 hMin(pos.x, pos.y + 1);
                                ImVec2 hMax(pos.x + hdrWidth, pos.y + blockHeight - 1);
                                dl->AddRectFilled(hMin, hMax, IM_COL32(40, 60, 120, 255));

                                if (hdrWidth > 30.0f)
                                {
                                    const char* hdrLabel = "Header";
                                    ImVec2      textSize = ImGui::CalcTextSize(hdrLabel);
                                    if (hdrWidth >= textSize.x + 4.0f)
                                    {
                                        float textX = hMin.x + (hdrWidth - textSize.x) * 0.5f;
                                        float textY = hMin.y + (blockHeight - 2.0f - textSize.y) * 0.5f;
                                        dl->AddText(ImVec2(textX, textY), IM_COL32(180, 200, 255, 255), hdrLabel);
                                    }
                                }

                                if (ImGui::IsMouseHoveringRect(hMin, hMax) && ImGui::IsWindowHovered())
                                {
                                    ImGui::BeginTooltip();
                                    ImGui::TextColored(ImVec4(0.5f, 0.6f, 0.9f, 1.0f), "Arena Header");
                                    ImGui::Separator();
                                    ImGui::Text("Size: %zu bytes", k_ArenaHeaderSize);
                                    ImGui::EndTooltip();
                                }
                            }

                            // Draw Allocations
                            ArenaDebugInfo* info           = blk->FirstDebugInfo;
                            size_t          expectedOffset = k_ArenaHeaderSize;
                            while (info)
                            {
                                // Draw alignment padding gap if present (only for small gaps that are actual alignment)
                                if (info->Offset > expectedOffset)
                                {
                                    size_t           padSize               = info->Offset - expectedOffset;
                                    constexpr size_t k_MaxAlignmentPadding = 256;
                                    float padXStart = pos.x + static_cast<float>(static_cast<double>(expectedOffset) * scale);
                                    float padWidth = static_cast<float>(static_cast<double>(padSize) * scale);

                                    if (padSize <= k_MaxAlignmentPadding && padWidth >= 1.0f)
                                    {
                                        ImVec2 pMin(padXStart, pos.y + 1);
                                        ImVec2 pMax(padXStart + padWidth, pos.y + blockHeight - 1);

                                        // Clip
                                        if (pMin.x < rectMin.x)
                                        {
                                            pMin.x = rectMin.x;
                                        }
                                        if (pMax.x > rectMax.x)
                                        {
                                            pMax.x = rectMax.x;
                                        }

                                        // Draw hatched pattern (diagonal lines) - clipped to visible area
                                        float visLeft      = ImGui::GetWindowPos().x;
                                        float visRight     = visLeft + ImGui::GetWindowSize().x;
                                        float clippedLeft  = std::max(pMin.x, visLeft);
                                        float clippedRight = std::min(pMax.x, visRight);

                                        dl->AddRectFilled(pMin, pMax, IM_COL32(30, 30, 30, 200));
                                        if (clippedLeft < clippedRight)
                                        {
                                            float step   = 6.0f;
                                            float startX = clippedLeft - (pMax.y - pMin.y);
                                            startX       = pMin.x + step * std::floor((startX - pMin.x) / step);
                                            if (startX < pMin.x)
                                            {
                                                startX = pMin.x;
                                            }

                                            int lineCount = 0;
                                            for (float lx = startX; lx < clippedRight && lineCount < 500; lx += step, ++lineCount)
                                            {
                                                float ly0 = pMin.y;
                                                float lx1 = lx + (pMax.y - pMin.y);
                                                float ly1 = pMax.y;
                                                if (lx1 > pMax.x)
                                                {
                                                    ly1 = pMin.y + (pMax.x - lx);
                                                    lx1 = pMax.x;
                                                }
                                                dl->AddLine(ImVec2(lx, ly0), ImVec2(lx1, ly1), IM_COL32(80, 80, 80, 180));
                                            }
                                        }

                                        // Label when zoomed
                                        if (padWidth > 30.0f)
                                        {
                                            const char* padLabel = "Pad";
                                            ImVec2      textSize = ImGui::CalcTextSize(padLabel);
                                            if (padWidth >= textSize.x + 4.0f)
                                            {
                                                float textX = pMin.x + (padWidth - textSize.x) * 0.5f;
                                                float textY = pMin.y + (blockHeight - 2.0f - textSize.y) * 0.5f;
                                                dl->AddText(ImVec2(textX, textY), IM_COL32(120, 120, 120, 255), padLabel);
                                            }
                                        }

                                        if (ImGui::IsMouseHoveringRect(pMin, pMax) && ImGui::IsWindowHovered())
                                        {
                                            ImGui::BeginTooltip();
                                            ImGui::TextColored(ImVec4(0.6f, 0.6f, 0.6f, 1.0f), "Alignment Padding");
                                            ImGui::Separator();
                                            ImGui::Text("Size: %zu bytes", padSize);
                                            ImGui::Text("Offset: %zu - %zu", expectedOffset, info->Offset);
                                            ImGui::EndTooltip();
                                        }
                                    }
                                }
                                expectedOffset = info->Offset + info->Size;

                                ImGui::PushID(info);
                                float xStart = pos.x + static_cast<float>(static_cast<double>(info->Offset) * scale);
                                float width  = static_cast<float>(static_cast<double>(info->Size) * scale);
                                width        = std::max(width, 1.0f);

                                ImVec2 aMin(xStart, pos.y + 1);
                                ImVec2 aMax(xStart + width, pos.y + blockHeight - 1);

                                // Clip to block bounds
                                if (aMin.x < rectMin.x)
                                {
                                    aMin.x = rectMin.x;
                                }
                                if (aMax.x > rectMax.x)
                                {
                                    aMax.x = rectMax.x;
                                }

                                ImU32 color      = GetColorForAllocation(info->Tag, info->Offset, info->Size);
                                bool  isSelected = (state.SelectedAlloc == info);

                                if (isSelected)
                                {
                                    dl->AddRectFilled(aMin, aMax, IM_COL32_WHITE);
                                    dl->AddRect(aMin, aMax, IM_COL32_BLACK);
                                }
                                else
                                {
                                    dl->AddRectFilled(aMin, aMax, color);
                                }

                                // Draw text label when zoomed enough
                                if (width > 20.0f && info->Tag)
                                {
                                    const char* tagStr = info->Tag;
                                    size_t      tagLen = 0;
                                    while (tagStr[tagLen] != '\0')
                                    {
                                        ++tagLen;
                                    }
                                    ImVec2 textSize = ImGui::CalcTextSize(tagStr, tagStr + tagLen);
                                    if (width >= textSize.x + 4.0f)
                                    {
                                        float textX     = aMin.x + (width - textSize.x) * 0.5f;
                                        float textY     = aMin.y + (blockHeight - 2.0f - textSize.y) * 0.5f;
                                        ImU32 textColor = IM_COL32_BLACK;
                                        dl->AddText(ImVec2(textX, textY), textColor, tagStr, tagStr + tagLen);
                                    }
                                }

                                if (ImGui::IsMouseHoveringRect(aMin, aMax) && ImGui::IsWindowHovered())
                                {
                                    ImGui::BeginTooltip();
                                    ImGui::TextColored(ImVec4(1.0f, 0.8f, 0.3f, 1.0f), "%s", info->Tag);
                                    ImGui::Separator();
                                    ImGui::Text("Size: %zu bytes", info->Size);
                                    ImGui::Text("Offset: %zu", info->Offset);

                                    // Hex data dump (first 16 bytes)
                                    ImGui::Separator();
                                    uint8* dataPtr = reinterpret_cast<uint8*>(const_cast<Arena*>(blk)) + info->Offset;
                                    ImGui::Text("Data: ");
                                    for (int i = 0; i < 16 && i < static_cast<int>(info->Size); ++i)
                                    {
                                        ImGui::SameLine();
                                        ImGui::Text("%02X", dataPtr[i]);
                                    }

                                    ImGui::EndTooltip();

                                    if (ImGui::IsMouseClicked(ImGuiMouseButton_Left))
                                    {
                                        state.SelectedAlloc          = info;
                                        state.ScrollListToSelected   = true;
                                        state.ScrollVisualToSelected = false;
                                    }
                                }

                                info = info->Next;
                                ImGui::PopID();
                            }

                            // Draw Free Nodes (Holes) if Realloc is enabled
                            if (arena->AllowRealloc && blk->FreeNodes)
                            {
                                ArenaFreeNode* freeNode = blk->FreeNodes;
                                while (freeNode)
                                {
                                    float fxStart = pos.x + static_cast<float>(static_cast<double>(freeNode->Position) * scale);
                                    float fWidth = static_cast<float>(static_cast<double>(freeNode->Size) * scale);

                                    ImVec2 fMin(fxStart, pos.y + 1);
                                    ImVec2 fMax(fxStart + fWidth, pos.y + blockHeight - 1);

                                    dl->AddRectFilled(fMin, fMax, IM_COL32(50, 50, 50, 200));

                                    freeNode = freeNode->Next;
                                }
                            }
                        }

                        pos.y += blockHeight + blockSpacing;
                        ImGui::PopID();
                    }
                }
                ImGui::EndChild();

                // --- Tree View (Allocations List) ---
                ImGui::Separator();
                if (ImGui::TreeNode("Allocations List"))
                {
                    // Calculate item count for dynamic height
                    int totalAllocCount = 0;
                    for (const Arena* blk : blocks)
                    {
                        totalAllocCount++;
                        ArenaDebugInfo* countInfo = blk->FirstDebugInfo;
                        while (countInfo)
                        {
                            totalAllocCount++;
                            countInfo = countInfo->Next;
                        }
                    }

                    float lineHeight          = ImGui::GetTextLineHeightWithSpacing();
                    float currentNeededHeight = static_cast<float>(totalAllocCount) * lineHeight;
                    float maxHeight           = lineHeight * 25.0f;

                    float listHeight = currentNeededHeight;
                    if (listHeight > maxHeight)
                    {
                        listHeight = maxHeight;
                    }
                    if (listHeight < lineHeight)
                    {
                        listHeight = lineHeight;
                    }

                    if (ImGui::BeginChild("AllocList", ImVec2(0, listHeight), true))
                    {
                        int blkIdx = 0;
                        for (const Arena* blk : blocks)
                        {
                            if (ImGui::TreeNode(reinterpret_cast<const void*>(blk), "Page %d (%zu bytes, pos %zu)",
                                                blkIdx, blk->Reserved, blk->Position))
                            {
                                ArenaDebugInfo* info = blk->FirstDebugInfo;
                                if (!info)
                                {
                                    ImGui::TextDisabled("No allocations");
                                }
                                while (info)
                                {
                                    bool isSelected = (state.SelectedAlloc == info);

                                    ImGuiTreeNodeFlags flags = ImGuiTreeNodeFlags_Leaf | ImGuiTreeNodeFlags_NoTreePushOnOpen;
                                    if (isSelected)
                                    {
                                        flags |= ImGuiTreeNodeFlags_Selected;
                                    }

                                    if (isSelected && state.ScrollListToSelected)
                                    {
                                        ImGui::SetScrollHereY();
                                        state.ScrollListToSelected = false;
                                    }

                                    // Color Square
                                    ImU32 color = GetColorForAllocation(info->Tag, info->Offset, info->Size);
                                    ImGui::ColorButton("##color", ImColor(color),
                                                       ImGuiColorEditFlags_NoTooltip | ImGuiColorEditFlags_NoDragDrop,
                                                       ImVec2(12, 12));
                                    ImGui::SameLine();

                                    ImGui::TreeNodeEx(info, flags, "[%zu] %s (%zu bytes)", info->Offset, info->Tag, info->Size);

                                    if (ImGui::IsItemClicked())
                                    {
                                        state.SelectedAlloc          = info;
                                        state.ScrollVisualToSelected = true;
                                        state.ScrollListToSelected   = false;
                                    }

                                    info = info->Next;
                                }
                                ImGui::TreePop();
                            }
                            blkIdx++;
                        }
                    }
                    ImGui::EndChild();

                    ImGui::TreePop();
                }
            }
            ImGui::PopID();
        }

        ArenaAllocation* s_ConfirmedHovered = nullptr;
    } // namespace

    void DebugDrawMemoryArena()
    {
        // Cross-frame hover state
        ArenaAllocation* frameHovered = nullptr;

        if (ImGui::Begin("Memory Debugger"))
        {
            if (ImGui::BeginTabBar("ArenaTabs"))
            {
                if (ImGui::BeginTabItem("Paged Arenas"))
                {
                    Arena* head = GetGlobalArenaListHead();
                    if (!head)
                    {
                        ImGui::Text("No active paged arenas.");
                    }
                    else
                    {
                        for (Arena* a = head; a != nullptr; a = a->GlobalNext)
                        {
                            // Skip test and internal debug arenas
                            if (a->Name)
                            {
                                bool isTest = (a->Name[0] == 'T' && a->Name[1] == 'e' && a->Name[2] == 's' && a->Name[3] == 't');
                                bool isDebugInfo =
                                    (a->Name[0] == 'D' && a->Name[1] == 'e' && a->Name[2] == 'b' && a->Name[3] == 'u' &&
                                     a->Name[4] == 'g' && a->Name[5] == ' ' && a->Name[6] == 'I');
                                if (isTest || isDebugInfo)
                                {
                                    continue;
                                }
                            }
                            DrawPagedArena(a);
                        }
                    }
                    ImGui::EndTabItem();
                }

                ImGui::EndTabBar();
            }
        }
        ImGui::End();

        s_ConfirmedHovered = frameHovered;
    }
} // namespace Juliet::Debug

#endif
