using ImGuiNET;
using System.Numerics;

namespace Nerfed.Runtime;

public static class ProfilerVisualizer
{
    private const float barHeight = 20f;
    private const float barPadding = 2f;

    // Render the flame graph across multiple threads
    public static void RenderFlameGraph(Profiler.Frame frame)
    {
        if (frame == null) return;
        if (frame.RootNodes == null) return;

        // Calculate the total timeline duration (max end time across all nodes)
        double totalDuration = frame.EndTime - frame.StartTime;
        double startTime = frame.StartTime;

        // Precompute the maximum depth for each thread's call stack
        Dictionary<int, int> threadMaxDepths = new Dictionary<int, int>();
        foreach (IGrouping<int, Profiler.ScopeNode> threadGroup in frame.RootNodes.GroupBy(node => node.ManagedThreadId))
        {
            int maxDepth = 0;
            foreach (Profiler.ScopeNode rootNode in threadGroup)
            {
                maxDepth = Math.Max(maxDepth, GetMaxDepth(rootNode, 0));
            }
            threadMaxDepths[threadGroup.Key] = maxDepth;
        }

        // Start a child window to support scrolling
        ImGui.BeginChild("FlameGraph", new Vector2(0, 64), ImGuiChildFlags.Border | ImGuiChildFlags.ResizeY, ImGuiWindowFlags.HorizontalScrollbar | ImGuiWindowFlags.AlwaysVerticalScrollbar);

        ImDrawListPtr drawList = ImGui.GetWindowDrawList();
        Vector2 windowPos = ImGui.GetCursorScreenPos();

        // Sort nodes by ThreadID, ensuring main thread (Thread ID 1) is on top
        IOrderedEnumerable<IGrouping<int, Profiler.ScopeNode>> threadGroups = frame.RootNodes.GroupBy(node => node.ManagedThreadId).OrderBy(g => g.Key);

        // Initial Y position for drawing
        float baseY = windowPos.Y;
        bool alternate = false;
        float contentWidth = ImGui.GetContentRegionAvail().X;

        // Draw each thread's flame graph row by row
        foreach (IGrouping<int, Profiler.ScopeNode> threadGroup in threadGroups)
        {
            int threadId = threadGroup.Key;

            // Compute the base Y position for this thread
            float threadBaseY = baseY;

            // Calculate the maximum height for this thread's flame graph
            float threadHeight = (threadMaxDepths[threadId] + 1) * (barHeight + barPadding);

            // Draw the alternating background for each thread row
            uint backgroundColor = ImGui.ColorConvertFloat4ToU32(alternate ? new Vector4(0.2f, 0.2f, 0.2f, 1f) : new Vector4(0.1f, 0.1f, 0.1f, 1f));
            drawList.AddRectFilled(new Vector2(windowPos.X, threadBaseY), new Vector2(windowPos.X + contentWidth, threadBaseY + threadHeight), backgroundColor);

            alternate = !alternate;

            // Draw each root node in the group (one per thread)
            foreach (Profiler.ScopeNode rootNode in threadGroup)
            {
                RenderNode(drawList, rootNode, startTime, totalDuration, windowPos.X, threadBaseY, 0, contentWidth, false);
            }

            // Move to the next thread's row (max depth * height per level)
            baseY += (threadMaxDepths[threadId] + 1) * (barHeight + barPadding);
        }

        // Ensure that ImGui knows the size of the content.
        ImGui.Dummy(new Vector2(contentWidth, baseY));

        ImGui.EndChild();
    }

    private static void RenderNode(ImDrawListPtr drawList, Profiler.ScopeNode node, double startTime, double totalDuration, float startX, float baseY, int depth, float contentWidth, bool alternate)
    {
        if (node == null) return;

        double nodeStartTime = node.StartTime - startTime;
        double nodeEndTime = node.EndTime - startTime;
        double nodeDuration = nodeEndTime - nodeStartTime;

        // Calculate the position and width of the bar based on time
        float xPos = (float)(startX + (nodeStartTime / totalDuration) * contentWidth);
        float width = (float)((nodeDuration / totalDuration) * contentWidth);

        // Calculate the Y position based on depth
        float yPos = baseY + (depth * (barHeight + barPadding)) + (barPadding * 0.5f);

        // Define the rectangle bounds for the node
        Vector2 min = new Vector2(xPos, yPos);
        Vector2 max = new Vector2(xPos + width, yPos + barHeight);

        // Define color.
        Vector4 barColor = alternate ? new Vector4(0.4f, 0.6f, 0.9f, 1f) : new Vector4(0.4f, 0.5f, 0.8f, 1f);
        Vector4 textColor = new Vector4(1f, 1f, 1f, 1f);

        if (depth != 0)
        {
            // Draw the bar for the node (colored based on thread depth)
            drawList.AddRectFilled(min, max, ImGui.ColorConvertFloat4ToU32(barColor));

            // Draw the label if it fits inside the bar
            string label = $"{node.Label} ({node.ElapsedMilliseconds():0.000} ms)";
            if (width > ImGui.CalcTextSize(label).X)
            {
                drawList.AddText(new Vector2(xPos + barPadding, yPos + barPadding), ImGui.ColorConvertFloat4ToU32(textColor), label);
            }

            // Add tooltip on hover
            if (ImGui.IsMouseHoveringRect(min, max))
            {
                // Show tooltip when hovering over the node
                ImGui.BeginTooltip();
                ImGui.Text($"{node.Label}");
                ImGui.Text($"{node.ElapsedMilliseconds():0.000} ms");
                ImGui.Text($"{node.ManagedThreadId}");
                ImGui.EndTooltip();
            }
        }
        else
        {
            // Aka root node.
            string label = $"{node.Label}";
            drawList.AddText(new Vector2(startX + barPadding, yPos + barPadding), ImGui.ColorConvertFloat4ToU32(textColor), label);
        }

        // Draw each child node under this node
        foreach (Profiler.ScopeNode child in node.Children)
        {
            alternate = !alternate;
            RenderNode(drawList, child, startTime, totalDuration, startX, baseY, depth + 1, contentWidth, alternate);
        }
    }

    // Recursive function to calculate the maximum depth of the node tree
    private static int GetMaxDepth(Profiler.ScopeNode node, int currentDepth)
    {
        if (node.Children == null || node.Children.Count == 0)
        {
            return currentDepth;
        }

        int maxDepth = currentDepth;
        foreach (Profiler.ScopeNode child in node.Children)
        {
            maxDepth = Math.Max(maxDepth, GetMaxDepth(child, currentDepth + 1));
        }
        return maxDepth;
    }
}