using ImGuiNET;
using System.Numerics;

namespace Nerfed.Runtime;

public static class ProfilerVisualizer
{
    public sealed class TimelineState
    {
        public int SelectedFrameIndex = -1;
        public int WindowStartIndex = 0;
        public int VisibleFrameCount = 64;
        public bool FollowLatest = true;
        public float Zoom = 1f;
        public double PanTicks = 0;
    }

    public readonly struct TimelineRenderResult
    {
        public TimelineRenderResult(int selectedFrameIndex, bool selectionChanged, bool userNavigated)
        {
            SelectedFrameIndex = selectedFrameIndex;
            SelectionChanged = selectionChanged;
            UserNavigated = userNavigated;
        }

        public int SelectedFrameIndex { get; }
        public bool SelectionChanged { get; }
        public bool UserNavigated { get; }
    }

    private readonly struct HoverEntry
    {
        public HoverEntry(Profiler.ScopeNode node, Profiler.Frame frame, int frameIndex, int depth, double timelineStartTicks)
        {
            Node = node;
            Frame = frame;
            FrameIndex = frameIndex;
            Depth = depth;
            DurationMs = TicksToMilliseconds(node.EndTime - node.StartTime);
            SelfMs = node.SelfMilliseconds();
            StartInFrameMs = TicksToMilliseconds(node.StartTime - frame.StartTime);
            EndInFrameMs = TicksToMilliseconds(node.EndTime - frame.StartTime);
            StartInTimelineMs = TicksToMilliseconds(node.StartTime - timelineStartTicks);
            EndInTimelineMs = TicksToMilliseconds(node.EndTime - timelineStartTicks);
        }

        public Profiler.ScopeNode Node { get; }
        public Profiler.Frame Frame { get; }
        public int FrameIndex { get; }
        public int Depth { get; }
        public double DurationMs { get; }
        public double SelfMs { get; }
        public double StartInFrameMs { get; }
        public double EndInFrameMs { get; }
        public double StartInTimelineMs { get; }
        public double EndInTimelineMs { get; }
    }

    private const float BarHeight = 18f;
    private const float BarPadding = 2f;
    private const float ThreadGap = 8f;
    private const float HeaderHeight = 28f;
    private const float TimelineHeight = 220f;
    private const float MinTextWidth = 36f;
    private static readonly double TickToMs = 1000d / System.Diagnostics.Stopwatch.Frequency;

    // Backwards-compatible entry point used by existing call sites.
    public static void RenderFlameGraph(Profiler.Frame frame)
    {
        if (frame == null)
        {
            return;
        }

        List<Profiler.Frame> frames = new List<Profiler.Frame>(1)
        {
            frame
        };

        TimelineState state = new TimelineState
        {
            VisibleFrameCount = 1,
            SelectedFrameIndex = 0,
            FollowLatest = true
        };

        RenderTimeline(frames, state);
    }

    public static TimelineRenderResult RenderTimeline(IReadOnlyList<Profiler.Frame> frames, TimelineState state)
    {
        if (frames == null || frames.Count == 0 || state == null)
        {
            return default;
        }

        bool selectionChanged = false;
        bool userNavigated = false;

        int frameCount = frames.Count;
        state.VisibleFrameCount = Math.Clamp(state.VisibleFrameCount, 1, frameCount);
        state.Zoom = Math.Clamp(state.Zoom, 1f, 128f);

        int maxStartIndex = Math.Max(0, frameCount - state.VisibleFrameCount);
        if (state.FollowLatest)
        {
            state.WindowStartIndex = maxStartIndex;
        }
        else
        {
            state.WindowStartIndex = Math.Clamp(state.WindowStartIndex, 0, maxStartIndex);
        }

        int visibleStartIndex = state.WindowStartIndex;
        int visibleEndIndex = visibleStartIndex + state.VisibleFrameCount - 1;

        if (state.SelectedFrameIndex < 0)
        {
            state.SelectedFrameIndex = visibleEndIndex;
            selectionChanged = true;
        }
        state.SelectedFrameIndex = Math.Clamp(state.SelectedFrameIndex, visibleStartIndex, visibleEndIndex);

        Profiler.Frame firstFrame = frames[visibleStartIndex];
        Profiler.Frame lastFrame = frames[visibleEndIndex];

        double timelineStartTicks = firstFrame.StartTime;
        double timelineEndTicks = Math.Max(lastFrame.EndTime, firstFrame.StartTime + 1);
        double timelineDurationTicks = Math.Max(1d, timelineEndTicks - timelineStartTicks);

        double visibleDurationTicks = Math.Max(1d, timelineDurationTicks / state.Zoom);
        double maxPanTicks = Math.Max(0d, timelineDurationTicks - visibleDurationTicks);
        if (state.FollowLatest)
        {
            state.PanTicks = maxPanTicks;
        }
        else
        {
            state.PanTicks = Math.Clamp(state.PanTicks, 0d, maxPanTicks);
        }

        double visibleStartTicks = timelineStartTicks + state.PanTicks;
        double visibleEndTicks = visibleStartTicks + visibleDurationTicks;

        Dictionary<int, int> threadDepths = BuildThreadDepths(frames, visibleStartIndex, visibleEndIndex);
        List<int> threadOrder = threadDepths.Keys.OrderBy(x => x).ToList();

        Dictionary<int, float> threadBaseY = new Dictionary<int, float>(threadOrder.Count);
        float yCursor = HeaderHeight;
        for (int i = 0; i < threadOrder.Count; i++)
        {
            int threadId = threadOrder[i];
            threadBaseY[threadId] = yCursor;
            yCursor += ((threadDepths[threadId] + 1) * (BarHeight + BarPadding)) + ThreadGap;
        }

        float contentHeight = Math.Max(TimelineHeight, yCursor + 6f);

        ImGui.BeginChild("ProfilerTimeline", new Vector2(0, TimelineHeight), ImGuiChildFlags.Border | ImGuiChildFlags.ResizeY, ImGuiWindowFlags.AlwaysVerticalScrollbar);

        ImDrawListPtr drawList = ImGui.GetWindowDrawList();
        Vector2 origin = ImGui.GetCursorScreenPos();
        Vector2 viewSize = ImGui.GetContentRegionAvail();
        float canvasWidth = Math.Max(1f, viewSize.X);

        uint frameBgColor = ImGui.GetColorU32(ImGuiCol.FrameBg);
        uint frameBgHoveredColor = ImGui.GetColorU32(ImGuiCol.FrameBgHovered);
        uint headerColor = ImGui.GetColorU32(ImGuiCol.Header);
        uint headerHoveredColor = ImGui.GetColorU32(ImGuiCol.HeaderHovered);
        uint textColor = ImGui.GetColorU32(ImGuiCol.Text);
        uint borderColor = ImGui.GetColorU32(ImGuiCol.Border);

        float clipMinX = origin.X;
        float clipMaxX = origin.X + canvasWidth;
        float clipMinY = origin.Y;
        float clipMaxY = origin.Y + Math.Max(1f, ImGui.GetWindowHeight());

        DrawTimelineHeader(drawList, origin, canvasWidth, timelineStartTicks, visibleStartTicks, visibleDurationTicks, textColor, borderColor);

        HoverEntry? hovered = null;

        for (int frameIndex = visibleStartIndex; frameIndex <= visibleEndIndex; frameIndex++)
        {
            Profiler.Frame frame = frames[frameIndex];
            float frameStartX = ToScreenX(frame.StartTime, visibleStartTicks, visibleDurationTicks, origin.X, canvasWidth);
            float frameEndX = ToScreenX(frame.EndTime, visibleStartTicks, visibleDurationTicks, origin.X, canvasWidth);

            if (frameEndX < clipMinX || frameStartX > clipMaxX)
            {
                continue;
            }

            bool isSelectedFrame = frameIndex == state.SelectedFrameIndex;
            uint frameShadeColor = (frameIndex & 1) == 0 ? frameBgColor : frameBgHoveredColor;
            if (isSelectedFrame)
            {
                frameShadeColor = LerpColor(headerColor, headerHoveredColor, 0.45f);
            }

            drawList.AddRectFilled(new Vector2(frameStartX, origin.Y + HeaderHeight), new Vector2(frameEndX, origin.Y + contentHeight), SetAlpha(frameShadeColor, isSelectedFrame ? 0.20f : 0.08f));
            drawList.AddLine(new Vector2(frameStartX, origin.Y + HeaderHeight), new Vector2(frameStartX, origin.Y + contentHeight), SetAlpha(borderColor, 0.55f), 1f);

            foreach (Profiler.ScopeNode root in frame.RootNodes)
            {
                if (!threadBaseY.TryGetValue(root.ManagedThreadId, out float baseY))
                {
                    continue;
                }

                DrawThreadLabel(drawList, origin.X, origin.Y + baseY, root.ManagedThreadId, textColor);
                for (int i = 0; i < root.Children.Count; i++)
                {
                    RenderNode(
                        drawList,
                        root.Children[i],
                        frame,
                        frameIndex,
                        baseY,
                        0,
                        visibleStartTicks,
                        visibleDurationTicks,
                        origin.X,
                        origin.Y,
                        canvasWidth,
                        clipMinX,
                        clipMaxX,
                        clipMaxY,
                        ref hovered,
                        textColor,
                        headerColor,
                        headerHoveredColor,
                        frameIndex == state.SelectedFrameIndex);
                }
            }

            drawList.AddLine(new Vector2(frameEndX, origin.Y + HeaderHeight), new Vector2(frameEndX, origin.Y + contentHeight), SetAlpha(borderColor, 0.30f), 1f);
        }

        ImGui.Dummy(new Vector2(canvasWidth, contentHeight));

        bool windowHovered = ImGui.IsWindowHovered(ImGuiHoveredFlags.ChildWindows);
        if (windowHovered)
        {
            HandleZoomAndPan(state, timelineStartTicks, timelineDurationTicks, visibleDurationTicks, origin.X, canvasWidth, ref userNavigated);
        }

        if (windowHovered && ImGui.IsMouseClicked(ImGuiMouseButton.Left))
        {
            int clickedFrame = FindFrameIndexByMouseX(frames, visibleStartIndex, visibleEndIndex, visibleStartTicks, visibleDurationTicks, origin.X, canvasWidth, ImGui.GetMousePos().X);
            if (clickedFrame >= visibleStartIndex && clickedFrame <= visibleEndIndex && clickedFrame != state.SelectedFrameIndex)
            {
                state.SelectedFrameIndex = clickedFrame;
                selectionChanged = true;
                state.FollowLatest = false;
                userNavigated = true;
            }
        }

        if (hovered.HasValue)
        {
            if (ImGui.IsMouseClicked(ImGuiMouseButton.Left) && hovered.Value.FrameIndex != state.SelectedFrameIndex)
            {
                state.SelectedFrameIndex = hovered.Value.FrameIndex;
                selectionChanged = true;
                state.FollowLatest = false;
                userNavigated = true;
            }

            DrawTooltip(hovered.Value);
        }

        ImGui.EndChild();

        return new TimelineRenderResult(state.SelectedFrameIndex, selectionChanged, userNavigated);
    }

    private static void HandleZoomAndPan(TimelineState state, double timelineStartTicks, double timelineDurationTicks, double visibleDurationTicks, float originX, float canvasWidth, ref bool userNavigated)
    {
        ImGuiIOPtr io = ImGui.GetIO();
        if (Math.Abs(io.MouseWheel) < float.Epsilon)
        {
            return;
        }

        bool isZoom = io.KeyCtrl;
        if (isZoom)
        {
            float previousZoom = state.Zoom;
            double visibleStartTicksBefore = timelineStartTicks + state.PanTicks;
            double mouseT = Math.Clamp((ImGui.GetMousePos().X - originX) / Math.Max(1f, canvasWidth), 0f, 1f);
            double pivotTick = visibleStartTicksBefore + (visibleDurationTicks * mouseT);

            state.Zoom = Math.Clamp(state.Zoom * MathF.Pow(1.12f, io.MouseWheel), 1f, 128f);
            if (Math.Abs(previousZoom - state.Zoom) > float.Epsilon)
            {
                double newVisibleDurationTicks = Math.Max(1d, timelineDurationTicks / state.Zoom);
                double newVisibleStartTicks = pivotTick - (newVisibleDurationTicks * mouseT);
                state.PanTicks = Math.Clamp(newVisibleStartTicks - timelineStartTicks, 0d, Math.Max(0d, timelineDurationTicks - newVisibleDurationTicks));
                state.FollowLatest = false;
                userNavigated = true;
            }
        }
        else
        {
            state.PanTicks -= io.MouseWheel * (visibleDurationTicks * 0.10d);
            state.FollowLatest = false;
            userNavigated = true;
        }
    }

    private static int FindFrameIndexByMouseX(IReadOnlyList<Profiler.Frame> frames, int visibleStartIndex, int visibleEndIndex, double visibleStartTicks, double visibleDurationTicks, float originX, float width, float mouseX)
    {
        double t = Math.Clamp((mouseX - originX) / Math.Max(1f, width), 0f, 1f);
        double timelineTicks = visibleStartTicks + (visibleDurationTicks * t);

        for (int i = visibleStartIndex; i <= visibleEndIndex; i++)
        {
            Profiler.Frame frame = frames[i];
            if (timelineTicks >= frame.StartTime && timelineTicks <= frame.EndTime)
            {
                return i;
            }
        }

        return -1;
    }

    private static void DrawTimelineHeader(ImDrawListPtr drawList, Vector2 origin, float width, double timelineStartTicks, double visibleStartTicks, double visibleDurationTicks, uint textColor, uint borderColor)
    {
        drawList.AddLine(new Vector2(origin.X, origin.Y + HeaderHeight), new Vector2(origin.X + width, origin.Y + HeaderHeight), SetAlpha(borderColor, 0.65f), 1f);

        int tickCount = Math.Clamp((int)(width / 130f), 4, 20);
        for (int i = 0; i <= tickCount; i++)
        {
            float t = i / (float)tickCount;
            float x = origin.X + (t * width);
            drawList.AddLine(new Vector2(x, origin.Y + HeaderHeight - 8f), new Vector2(x, origin.Y + HeaderHeight), SetAlpha(borderColor, 0.8f), 1f);

            double ms = TicksToMilliseconds((visibleStartTicks - timelineStartTicks) + (visibleDurationTicks * t));
            drawList.AddText(new Vector2(x + 2f, origin.Y + 4f), textColor, $"+{ms:0.0} ms");
        }
    }

    private static void DrawThreadLabel(ImDrawListPtr drawList, float x, float y, int threadId, uint textColor)
    {
        drawList.AddText(new Vector2(x + 4f, y + 2f), SetAlpha(textColor, 0.85f), $"T{threadId}");
    }

    private static void DrawTooltip(HoverEntry hover)
    {
        ImGui.BeginTooltip();
        ImGui.Text($"{hover.Node.Label}");
        ImGui.Separator();
        ImGui.Text($"Category: {hover.Node.Category}");
        ImGui.Text($"Tags: 0x{hover.Node.TagMask:X}");
        ImGui.Text($"Frame: {hover.Frame.FrameCount} (idx {hover.FrameIndex})");
        ImGui.Text($"Thread: {hover.Node.ManagedThreadId}");
        ImGui.Text($"Depth: {hover.Depth}");
        ImGui.Text($"Duration: {hover.DurationMs:0.000} ms");
        ImGui.Text($"Self: {hover.SelfMs:0.000} ms");
        ImGui.Text($"Frame Start: {hover.StartInFrameMs:0.000} ms");
        ImGui.Text($"Frame End: {hover.EndInFrameMs:0.000} ms");
        ImGui.Text($"Timeline Start: {hover.StartInTimelineMs:0.000} ms");
        ImGui.Text($"Timeline End: {hover.EndInTimelineMs:0.000} ms");
        ImGui.Text($"Children: {hover.Node.Children.Count}");
        ImGui.EndTooltip();
    }

    private static void RenderNode(
        ImDrawListPtr drawList,
        Profiler.ScopeNode node,
        Profiler.Frame frame,
        int frameIndex,
        float baseY,
        int depth,
        double visibleStartTicks,
        double visibleDurationTicks,
        float originX,
        float originY,
        float width,
        float clipMinX,
        float clipMaxX,
        float clipMaxY,
        ref HoverEntry? hovered,
        uint textColor,
        uint headerColor,
        uint headerHoveredColor,
        bool selectedFrame)
    {
        long nodeEndTime = Math.Max(node.EndTime, node.StartTime + 1);
        if (nodeEndTime < visibleStartTicks || node.StartTime > visibleStartTicks + visibleDurationTicks)
        {
            return;
        }

        float y = originY + baseY + (depth * (BarHeight + BarPadding));
        if (y > clipMaxY)
        {
            return;
        }

        float minX = ToScreenX(node.StartTime, visibleStartTicks, visibleDurationTicks, originX, width);
        float maxX = ToScreenX(nodeEndTime, visibleStartTicks, visibleDurationTicks, originX, width);
        if (maxX < clipMinX || minX > clipMaxX)
        {
            return;
        }

        float barWidth = Math.Max(1f, maxX - minX);
        Vector2 min = new Vector2(minX, y + (BarPadding * 0.5f));
        Vector2 max = new Vector2(minX + barWidth, y + (BarPadding * 0.5f) + BarHeight);

        uint barColor = BuildBarColor(node.Label, depth, selectedFrame);
        uint borderColor = LerpColor(headerColor, headerHoveredColor, 0.45f);
        drawList.AddRectFilled(min, max, barColor, 3f);
        drawList.AddRect(min, max, SetAlpha(borderColor, 0.55f), 3f, ImDrawFlags.None, 1f);

        Vector2 mousePos = ImGui.GetMousePos();
        bool isHovered = mousePos.X >= min.X && mousePos.X <= max.X && mousePos.Y >= min.Y && mousePos.Y <= max.Y;
        if (isHovered)
        {
            hovered = new HoverEntry(node, frame, frameIndex, depth, visibleStartTicks);
            drawList.AddRect(min, max, headerHoveredColor, 3f, ImDrawFlags.None, 1.5f);
        }

        if (barWidth > MinTextWidth)
        {
            string label = node.Label;
            float textWidth = ImGui.CalcTextSize(label).X;
            if (textWidth + 8f <= barWidth)
            {
                drawList.AddText(new Vector2(min.X + 4f, min.Y + 2f), textColor, label);
            }
        }

        for (int i = 0; i < node.Children.Count; i++)
        {
            RenderNode(
                drawList,
                node.Children[i],
                frame,
                frameIndex,
                baseY,
                depth + 1,
                visibleStartTicks,
                visibleDurationTicks,
                originX,
                originY,
                width,
                clipMinX,
                clipMaxX,
                clipMaxY,
                ref hovered,
                textColor,
                headerColor,
                headerHoveredColor,
                selectedFrame);
        }
    }

    private static Dictionary<int, int> BuildThreadDepths(IReadOnlyList<Profiler.Frame> frames, int startIndex, int endIndex)
    {
        Dictionary<int, int> threadMaxDepths = new Dictionary<int, int>(8);
        for (int frameIndex = startIndex; frameIndex <= endIndex; frameIndex++)
        {
            foreach (Profiler.ScopeNode root in frames[frameIndex].RootNodes)
            {
                int maxDepth = 0;
                for (int i = 0; i < root.Children.Count; i++)
                {
                    maxDepth = Math.Max(maxDepth, GetMaxDepth(root.Children[i], 0));
                }

                if (threadMaxDepths.TryGetValue(root.ManagedThreadId, out int currentMax))
                {
                    if (maxDepth > currentMax)
                    {
                        threadMaxDepths[root.ManagedThreadId] = maxDepth;
                    }
                }
                else
                {
                    threadMaxDepths[root.ManagedThreadId] = maxDepth;
                }
            }
        }

        return threadMaxDepths;
    }

    private static int GetMaxDepth(Profiler.ScopeNode node, int depth)
    {
        if (node.Children.Count == 0)
        {
            return depth;
        }

        int maxDepth = depth;
        for (int i = 0; i < node.Children.Count; i++)
        {
            maxDepth = Math.Max(maxDepth, GetMaxDepth(node.Children[i], depth + 1));
        }
        return maxDepth;
    }

    private static uint BuildBarColor(string label, int depth, bool selectedFrame)
    {
        int hash = label.GetHashCode();
        float hue = ((hash & 1023) / 1023f + (depth * 0.031f)) % 1f;
        ImGui.ColorConvertHSVtoRGB(hue, 0.52f, selectedFrame ? 0.82f : 0.68f, out float r, out float g, out float b);

        Vector4 frameBg = ImGui.ColorConvertU32ToFloat4(ImGui.GetColorU32(ImGuiCol.FrameBg));
        Vector4 accent = new Vector4(r, g, b, 1f);
        Vector4 mixed = Vector4.Lerp(frameBg, accent, 0.72f);
        return ImGui.ColorConvertFloat4ToU32(mixed);
    }

    private static float ToScreenX(double ticks, double visibleStartTicks, double visibleDurationTicks, float startX, float width)
    {
        double normalized = (ticks - visibleStartTicks) / visibleDurationTicks;
        return startX + (float)(normalized * width);
    }

    private static double TicksToMilliseconds(double ticks)
    {
        return ticks * TickToMs;
    }

    private static uint SetAlpha(uint color, float alpha)
    {
        Vector4 c = ImGui.ColorConvertU32ToFloat4(color);
        c.W *= alpha;
        return ImGui.ColorConvertFloat4ToU32(c);
    }

    private static uint LerpColor(uint a, uint b, float t)
    {
        Vector4 av = ImGui.ColorConvertU32ToFloat4(a);
        Vector4 bv = ImGui.ColorConvertU32ToFloat4(b);
        return ImGui.ColorConvertFloat4ToU32(Vector4.Lerp(av, bv, Math.Clamp(t, 0f, 1f)));
    }
}