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

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, double selfMs, uint calls, double avgMs, double p95Ms, long allocBytes, long selfAllocBytes)>> orderedCombinedData = null;
        private IOrderedEnumerable<KeyValuePair<int, Profiler.RollingThreadMetrics>> orderedThreadRollingData = null;
        private readonly ProfilerVisualizer.TimelineState timelineState = new ProfilerVisualizer.TimelineState();
        private readonly List<Profiler.Frame> frameSnapshot = new List<Profiler.Frame>(256);

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

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

            if (Profiler.CopyFramesTo(frameSnapshot) <= 0)
            {
                return;
            }

            selectedFrame = Math.Clamp(selectedFrame, 0, frameSnapshot.Count - 1);
            timelineState.SelectedFrameIndex = Math.Clamp(timelineState.SelectedFrameIndex, -1, frameSnapshot.Count - 1);
            timelineState.VisibleFrameCount = Math.Clamp(timelineState.VisibleFrameCount, 1, frameSnapshot.Count);

            ImGui.Begin("Profiler");

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

            ImGui.SameLine();
            ImGui.Text("Follow");
            ImGui.SameLine();
            ImGui.Checkbox("##follow-timeline", ref timelineState.FollowLatest);
            ImGui.SameLine();

            int visibleFrameCount = timelineState.VisibleFrameCount;
            ImGui.SetNextItemWidth(130f);
            if (ImGui.SliderInt("Window", ref visibleFrameCount, 1, frameSnapshot.Count))
            {
                timelineState.VisibleFrameCount = visibleFrameCount;
                timelineState.FollowLatest = false;
            }

            ImGui.SameLine();
            if (ImGui.Button("Reset Zoom"))
            {
                timelineState.Zoom = 1f;
                timelineState.PanTicks = 0d;
                timelineState.FollowLatest = true;
            }

            ImGui.SameLine();
            int mode = (int)Profiler.Mode;
            ImGui.SetNextItemWidth(130f);
            if (ImGui.Combo("Mode", ref mode, "Instrumented\0Sampled\0"))
            {
                Profiler.Mode = (Profiler.CaptureMode)mode;
            }

            ImGui.SameLine();
            int stride = Profiler.SamplingStride;
            ImGui.SetNextItemWidth(90f);
            if (ImGui.SliderInt("Stride", ref stride, 1, 64))
            {
                Profiler.SamplingStride = stride;
            }

            if (Profiler.IsRecording)
            {
                // Select last frame when recording to see latest frame data.
                selectedFrame = frameSnapshot.Count - 1;
            }

            int sliderFrame = selectedFrame;
            if (ImGui.SliderInt("Frame", ref sliderFrame, 0, frameSnapshot.Count - 1))
            {
                selectedFrame = sliderFrame;
                timelineState.SelectedFrameIndex = selectedFrame;
                timelineState.FollowLatest = false;
            }

            Profiler.Frame frame = frameSnapshot[selectedFrame];
            double ms = frame.ElapsedMilliseconds();
            double s = 1000;
            ImGui.Text($"Frame: {frame.FrameCount} ({ms:0.000} ms | {(s / ms):0} fps)");
            ImGui.Text($"Budget: {frame.BudgetMilliseconds:0.00} ms ({(frame.OverBudget ? "over" : "within")})");
            ImGui.Text($"Thread Budget: {Profiler.ThreadBudgetMilliseconds:0.00} ms | Capture: {Profiler.Mode}");
            ImGui.Text($"Alloc: {frame.AllocatedBytesDelta / 1024d:0.0} KB | GC: G0 {frame.Gen0CollectionsDelta}, G1 {frame.Gen1CollectionsDelta}, G2 {frame.Gen2CollectionsDelta}");
            ImGui.EndChild();

            ProfilerVisualizer.TimelineRenderResult timelineResult = DrawFlameGraph(frameSnapshot, timelineState);
            if (timelineResult.SelectionChanged)
            {
                selectedFrame = timelineResult.SelectedFrameIndex;
            }

            selectedFrame = Math.Clamp(selectedFrame, 0, frameSnapshot.Count - 1);
            frame = frameSnapshot[selectedFrame];

            if (previousSelectedFrame != selectedFrame)
            {
                previousSelectedFrame = selectedFrame;
                orderedCombinedData = CalculateCombinedData(frame);
                orderedThreadRollingData = CalculateThreadRollingData();
            }

            DrawThreadRolling(orderedThreadRollingData);

            DrawHierachy(frame);

            ImGui.SameLine();

            DrawCombined(orderedCombinedData);

            ImGui.End();
        }

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

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

            if (ImGui.BeginTable("ProfilerData", 8, tableFlags, new Vector2(0, 0)))
            {
                ImGui.TableSetupColumn("name", ImGuiTableColumnFlags.WidthStretch, 0.30f, 0);
                ImGui.TableSetupColumn("category", ImGuiTableColumnFlags.WidthStretch, 0.12f, 1);
                ImGui.TableSetupColumn("tags", ImGuiTableColumnFlags.WidthStretch, 0.08f, 2);
                ImGui.TableSetupColumn("thread", ImGuiTableColumnFlags.WidthStretch, 0.08f, 3);
                ImGui.TableSetupColumn("ms", ImGuiTableColumnFlags.WidthStretch, 0.10f, 4);
                ImGui.TableSetupColumn("self", ImGuiTableColumnFlags.WidthStretch, 0.10f, 5);
                ImGui.TableSetupColumn("alloc(B)", ImGuiTableColumnFlags.WidthStretch, 0.11f, 6);
                ImGui.TableSetupColumn("self alloc", ImGuiTableColumnFlags.WidthStretch, 0.11f, 7);
                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.Category}");
            ImGui.TableNextColumn();
            ImGui.Text($"0x{node.TagMask:X}");
            ImGui.TableNextColumn();
            ImGui.Text($"{node.ManagedThreadId}");
            ImGui.TableNextColumn();
            ImGui.Text($"{node.ElapsedMilliseconds():0.000}");
            ImGui.TableNextColumn();
            ImGui.Text($"{node.SelfMilliseconds():0.000}");
            ImGui.TableNextColumn();
            if (node.ProfilerSetupBytes > 0)
            {
                ImGui.Text($"{node.AllocatedBytes} !");
                if (ImGui.IsItemHovered())
                    ImGui.SetTooltip($"{node.ProfilerSetupBytes} B is profiler warmup overhead");
            }
            else
            {
                ImGui.Text($"{node.AllocatedBytes}");
            }
            ImGui.TableNextColumn();
            ImGui.Text($"{node.SelfAllocatedBytes()}");

            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, double selfMs, uint calls, double avgMs, double p95Ms, long allocBytes, long selfAllocBytes)>> orderedCombinedData)
        {
            if(orderedCombinedData == null)
            {
                return;
            }

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

            if (ImGui.BeginTable("ProfilerCombinedData", 8, tableFlags, new Vector2(0, 0)))
            {
                ImGui.TableSetupColumn("name", ImGuiTableColumnFlags.WidthStretch, 0.32f, 0);
                ImGui.TableSetupColumn("ms", ImGuiTableColumnFlags.WidthStretch, 0.11f, 1);
                ImGui.TableSetupColumn("self", ImGuiTableColumnFlags.WidthStretch, 0.11f, 2);
                ImGui.TableSetupColumn("calls", ImGuiTableColumnFlags.WidthStretch, 0.08f, 3);
                ImGui.TableSetupColumn("avg", ImGuiTableColumnFlags.WidthStretch, 0.09f, 4);
                ImGui.TableSetupColumn("p95", ImGuiTableColumnFlags.WidthStretch, 0.09f, 5);
                ImGui.TableSetupColumn("alloc(B)", ImGuiTableColumnFlags.WidthStretch, 0.10f, 6);
                ImGui.TableSetupColumn("self alloc", ImGuiTableColumnFlags.WidthStretch, 0.10f, 7);
                ImGui.TableSetupScrollFreeze(0, 1); // Make row always visible
                ImGui.TableHeadersRow();

                foreach (KeyValuePair<string, (double ms, double selfMs, uint calls, double avgMs, double p95Ms, long allocBytes, long selfAllocBytes)> 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.selfMs:0.000}");
                    ImGui.TableNextColumn();
                    ImGui.Text($"{combinedData.Value.calls}");
                    ImGui.TableNextColumn();
                    ImGui.Text($"{combinedData.Value.avgMs:0.000}");
                    ImGui.TableNextColumn();
                    ImGui.Text($"{combinedData.Value.p95Ms:0.000}");
                    ImGui.TableNextColumn();
                    ImGui.Text($"{combinedData.Value.allocBytes}");
                    ImGui.TableNextColumn();
                    ImGui.Text($"{combinedData.Value.selfAllocBytes}");
                }

                ImGui.EndTable();
            }

            ImGui.EndChild();
        }

        private static IOrderedEnumerable<KeyValuePair<string, (double ms, double selfMs, uint calls, double avgMs, double p95Ms, long allocBytes, long selfAllocBytes)>> CalculateCombinedData(Profiler.Frame frame)
        {
            IReadOnlyDictionary<string, Profiler.RollingLabelMetrics> rollingData = Profiler.GetRollingLabelMetricsSnapshot();
            Dictionary<string, (double ms, double selfMs, uint calls, double avgMs, double p95Ms, long allocBytes, long selfAllocBytes)> combinedRecordData = new Dictionary<string, (double ms, double selfMs, uint calls, double avgMs, double p95Ms, long allocBytes, long selfAllocBytes)>(128);
            foreach (KeyValuePair<string, Profiler.LabelMetrics> metric in frame.LabelMetrics)
            {
                Profiler.RollingLabelMetrics rolling = default;
                if (rollingData.TryGetValue(metric.Key, out Profiler.RollingLabelMetrics found))
                {
                    rolling = found;
                }

                combinedRecordData[metric.Key] = (metric.Value.InclusiveMs, metric.Value.SelfMs, metric.Value.Calls, rolling.AverageMs, rolling.P95Ms, metric.Value.AllocatedBytes, metric.Value.SelfAllocatedBytes);
            }
            return combinedRecordData.OrderByDescending(x => x.Value.ms);
        }

        private static IOrderedEnumerable<KeyValuePair<int, Profiler.RollingThreadMetrics>> CalculateThreadRollingData()
        {
            IReadOnlyDictionary<int, Profiler.RollingThreadMetrics> rollingData = Profiler.GetRollingThreadMetricsSnapshot();
            return rollingData.OrderByDescending(x => x.Value.P95Ms);
        }

        private static void DrawThreadRolling(in IOrderedEnumerable<KeyValuePair<int, Profiler.RollingThreadMetrics>> orderedThreadRollingData)
        {
            if (orderedThreadRollingData == null)
            {
                return;
            }

            ImGui.BeginChild("ThreadRolling", new Vector2(0, 140), ImGuiChildFlags.Border);
            if (ImGui.BeginTable("ProfilerThreadRollingData", 6, tableFlags, new Vector2(0, 0)))
            {
                ImGui.TableSetupColumn("thread", ImGuiTableColumnFlags.WidthStretch, 0.15f, 0);
                ImGui.TableSetupColumn("avg", ImGuiTableColumnFlags.WidthStretch, 0.20f, 1);
                ImGui.TableSetupColumn("p95", ImGuiTableColumnFlags.WidthStretch, 0.20f, 2);
                ImGui.TableSetupColumn("max", ImGuiTableColumnFlags.WidthStretch, 0.20f, 3);
                ImGui.TableSetupColumn("samples", ImGuiTableColumnFlags.WidthStretch, 0.15f, 4);
                ImGui.TableSetupColumn("misses", ImGuiTableColumnFlags.WidthStretch, 0.15f, 5);
                ImGui.TableHeadersRow();

                foreach (KeyValuePair<int, Profiler.RollingThreadMetrics> metric in orderedThreadRollingData)
                {
                    ImGui.TableNextRow();
                    ImGui.TableNextColumn();
                    ImGui.Text($"T{metric.Key}");
                    ImGui.TableNextColumn();
                    ImGui.Text($"{metric.Value.AverageMs:0.000}");
                    ImGui.TableNextColumn();
                    ImGui.Text($"{metric.Value.P95Ms:0.000}");
                    ImGui.TableNextColumn();
                    ImGui.Text($"{metric.Value.MaxMs:0.000}");
                    ImGui.TableNextColumn();
                    ImGui.Text($"{metric.Value.Samples}");
                    ImGui.TableNextColumn();
                    ImGui.Text($"{metric.Value.BudgetMisses}");
                }

                ImGui.EndTable();
            }

            ImGui.EndChild();
        }

        private static ProfilerVisualizer.TimelineRenderResult DrawFlameGraph(IReadOnlyList<Profiler.Frame> frames, ProfilerVisualizer.TimelineState timelineState)
        {
            if (frames == null || frames.Count == 0)
            {
                return default;
            }

            return ProfilerVisualizer.RenderTimeline(frames, timelineState);
        }
    }
}