using ImGuiNET;
using MoonTools.ECS;
using Nerfed.Runtime;

namespace Nerfed.Editor.Systems
{
    internal class EditorProfilerWindow : MoonTools.ECS.System
    {
        const ImGuiTableFlags tableFlags = ImGuiTableFlags.Resizable | ImGuiTableFlags.BordersOuter | ImGuiTableFlags.NoBordersInBody | ImGuiTableFlags.ScrollY | ImGuiTableFlags.ScrollX;
        const ImGuiTreeNodeFlags treeNodeFlags = ImGuiTreeNodeFlags.SpanAllColumns;
        const ImGuiTreeNodeFlags treeNodeLeafFlags = ImGuiTreeNodeFlags.SpanAllColumns | ImGuiTreeNodeFlags.Leaf | ImGuiTreeNodeFlags.NoTreePushOnOpen;

        private int selectedFrame = 0;
        private int previousSelectedFrame = -1;
        private IOrderedEnumerable<KeyValuePair<string, (double ms, uint calls)>> orderedCombinedData = null;

        public EditorProfilerWindow(World world) : base(world)
        {
        }

        public override void Update(TimeSpan delta)
        {
            if (Profiler.Frames.Count <= 0)
            {
                return;
            }

            ImGui.Begin("Profiler");

            ImGui.BeginChild("Toolbar", new System.Numerics.Vector2(0, 0), ImGuiChildFlags.AutoResizeY);
            if (ImGui.RadioButton("Recording", Profiler.IsRecording))
            {
                Profiler.SetActive(!Profiler.IsRecording);
            }

            ImGui.SameLine();

            if (Profiler.IsRecording)
            {
                // Select last frame when recording to see latest frame data.
                selectedFrame = Profiler.Frames.Count - 1;
            }
            if (ImGui.SliderInt(string.Empty, ref selectedFrame, 0, Profiler.Frames.Count - 1))
            {
                // Stop recording when browsing frames.
                Profiler.SetActive(false);
            }

            Profiler.Frame frame = Profiler.Frames.ElementAt(selectedFrame);
            double ms = frame.ElapsedMilliseconds();
            double s = 1000;
            ImGui.Text($"Frame: {frame.FrameCount} ({ms:0.000} ms | {(s / ms):0} fps)");
            ImGui.EndChild();

            if (!Profiler.IsRecording) {
                if (previousSelectedFrame != selectedFrame)
                {
                    previousSelectedFrame = selectedFrame;
                    orderedCombinedData = CalculateCombinedData(frame);
                }

                DrawFlameGraph(frame);

                DrawHierachy(frame);

                ImGui.SameLine();

                DrawCombined(orderedCombinedData);
            }

            ImGui.End();
        }

        private static void DrawHierachy(Profiler.Frame frame)
        {
            if(frame == null)
            {
                return;
            }

            ImGui.BeginChild("Hierachy", new System.Numerics.Vector2(150, 0), ImGuiChildFlags.ResizeX);

            if (ImGui.BeginTable("ProfilerData", 3, tableFlags, new System.Numerics.Vector2(0, 0)))
            {
                ImGui.TableSetupColumn("name", ImGuiTableColumnFlags.WidthStretch, 0.8f, 0);
                ImGui.TableSetupColumn("thread", ImGuiTableColumnFlags.WidthStretch, 0.2f, 1);
                ImGui.TableSetupColumn("ms", ImGuiTableColumnFlags.WidthStretch, 0.2f, 1);
                ImGui.TableSetupScrollFreeze(0, 1); // Make row always visible
                ImGui.TableHeadersRow();

                foreach (Profiler.ScopeNode node in frame.RootNodes)
                {
                    DrawHierachyNode(node);
                }

                ImGui.EndTable();
            }

            ImGui.EndChild();
        }

        private static void DrawHierachyNode(Profiler.ScopeNode node)
        {
            ImGui.TableNextRow();
            ImGui.TableNextColumn();

            bool isOpen = false;
            bool isLeaf = node.Children.Count == 0;

            if (isLeaf) {
                ImGui.TreeNodeEx(node.Label, treeNodeLeafFlags);
            }
            else
            {
                isOpen = ImGui.TreeNodeEx(node.Label, treeNodeFlags);
            }

            ImGui.TableNextColumn();
            ImGui.Text($"{node.ManagedThreadId}");
            ImGui.TableNextColumn();
            ImGui.Text($"{node.ElapsedMilliseconds():0.000}");

            if (isOpen)
            {
                for (int i = 0; i < node.Children.Count; i++)
                {
                    DrawHierachyNode(node.Children[i]);
                }
                ImGui.TreePop();
            }
        }

        private static void DrawCombined(in IOrderedEnumerable<KeyValuePair<string, (double ms, uint calls)>> orderedCombinedData)
        {
            if(orderedCombinedData == null)
            {
                return;
            }

            ImGui.BeginChild("Combined", new System.Numerics.Vector2(0, 0));

            if (ImGui.BeginTable("ProfilerCombinedData", 3, tableFlags, new System.Numerics.Vector2(0, 0)))
            {
                ImGui.TableSetupColumn("name", ImGuiTableColumnFlags.WidthStretch, 0.6f, 0);
                ImGui.TableSetupColumn("ms", ImGuiTableColumnFlags.WidthStretch, 0.2f, 1);
                ImGui.TableSetupColumn("calls", ImGuiTableColumnFlags.WidthStretch, 0.2f, 2);
                ImGui.TableSetupScrollFreeze(0, 1); // Make row always visible
                ImGui.TableHeadersRow();

                foreach (KeyValuePair<string, (double ms, uint calls)> combinedData in orderedCombinedData)
                {
                    ImGui.TableNextRow();
                    ImGui.TableNextColumn();
                    ImGui.Text($"{combinedData.Key}");
                    ImGui.TableNextColumn();
                    ImGui.Text($"{combinedData.Value.ms:0.000}");
                    ImGui.TableNextColumn();
                    ImGui.Text($"{combinedData.Value.calls}");
                }

                ImGui.EndTable();
            }

            ImGui.EndChild();
        }

        private static IOrderedEnumerable<KeyValuePair<string, (double ms, uint calls)>> CalculateCombinedData(Profiler.Frame frame)
        {
            Dictionary<string, (double ms, uint calls)> combinedRecordData = new Dictionary<string, (double ms, uint calls)>(128);
            foreach (Profiler.ScopeNode node in frame.RootNodes)
            {
                CalculateCombinedData(node, in combinedRecordData);
            }
            return combinedRecordData.OrderByDescending(x => x.Value.ms);
        }

        private static void CalculateCombinedData(Profiler.ScopeNode node, in Dictionary<string, (double ms, uint calls)> combinedRecordData)
        {
            if (combinedRecordData.TryGetValue(node.Label, out (double ms, uint calls) combined))
            {
                combinedRecordData[node.Label] = (combined.ms + node.ElapsedMilliseconds(), combined.calls + 1);
            }
            else
            {
                combinedRecordData.Add(node.Label, (node.ElapsedMilliseconds(), 1));
            }

            for (int i = 0; i < node.Children.Count; i++)
            {
                CalculateCombinedData(node.Children[i], combinedRecordData);
            }
        }

        private static void DrawFlameGraph(Profiler.Frame frame)
        {
            if (frame == null)
            {
                return;
            }

            ProfilerVisualizer.RenderFlameGraph(frame);
        }
    }
}