using System.Collections.Concurrent;
using System.Diagnostics;

namespace Nerfed.Runtime;

public struct ProfilerScope : IDisposable
{
    public ProfilerScope(string label)
    {
        Profiler.BeginSample(label);
    }

    public void Dispose()
    {
        Profiler.EndSample();
    }
}

public static class Profiler
{
    public readonly struct LabelMetrics
    {
        public LabelMetrics(double inclusiveMs, double selfMs, uint calls, double minInclusiveMs, double maxInclusiveMs)
        {
            InclusiveMs = inclusiveMs;
            SelfMs = selfMs;
            Calls = calls;
            MinInclusiveMs = minInclusiveMs;
            MaxInclusiveMs = maxInclusiveMs;
        }

        public double InclusiveMs { get; }
        public double SelfMs { get; }
        public uint Calls { get; }
        public double MinInclusiveMs { get; }
        public double MaxInclusiveMs { get; }
    }

    public readonly struct RollingLabelMetrics
    {
        public RollingLabelMetrics(double averageMs, double minMs, double maxMs, double p95Ms, int samples)
        {
            AverageMs = averageMs;
            MinMs = minMs;
            MaxMs = maxMs;
            P95Ms = p95Ms;
            Samples = samples;
        }

        public double AverageMs { get; }
        public double MinMs { get; }
        public double MaxMs { get; }
        public double P95Ms { get; }
        public int Samples { get; }
    }

    private sealed class RollingWindow
    {
        private readonly double[] values;
        private int index;
        private int count;

        public RollingWindow(int capacity)
        {
            values = new double[Math.Max(8, capacity)];
        }

        public void Add(double value)
        {
            values[index] = value;
            index = (index + 1) % values.Length;
            if (count < values.Length)
            {
                count++;
            }
        }

        public RollingLabelMetrics Snapshot()
        {
            if (count == 0)
            {
                return default;
            }

            double sum = 0;
            double min = double.MaxValue;
            double max = double.MinValue;
            double[] sorted = new double[count];

            int start = (index - count + values.Length) % values.Length;
            for (int i = 0; i < count; i++)
            {
                double value = values[(start + i) % values.Length];
                sorted[i] = value;
                sum += value;
                min = Math.Min(min, value);
                max = Math.Max(max, value);
            }

            Array.Sort(sorted);
            int percentileIndex = (int)Math.Ceiling((count - 1) * 0.95d);
            double p95 = sorted[Math.Clamp(percentileIndex, 0, count - 1)];
            return new RollingLabelMetrics(sum / count, min, max, p95, count);
        }
    }

    public class Frame(uint frameCount)
    {
        public uint FrameCount { get; } = frameCount;
        public long StartTime { get; } = Stopwatch.GetTimestamp();
        public long EndTime { get; private set; }

        private readonly List<ScopeNode> rootNodes = new List<ScopeNode>(8);
        private readonly object rootNodesLock = new object();
        private readonly Dictionary<string, LabelMetrics> labelMetrics = new Dictionary<string, LabelMetrics>(128, StringComparer.Ordinal);

        public IReadOnlyList<ScopeNode> RootNodes => rootNodes;

        public IReadOnlyDictionary<string, LabelMetrics> LabelMetrics => labelMetrics;
        public long AllocatedBytesStart { get; } = GC.GetTotalAllocatedBytes(false);
        public long AllocatedBytesEnd { get; private set; }
        public long AllocatedBytesDelta { get; private set; }
        public int Gen0CollectionsStart { get; } = GC.CollectionCount(0);
        public int Gen1CollectionsStart { get; } = GC.CollectionCount(1);
        public int Gen2CollectionsStart { get; } = GC.CollectionCount(2);
        public int Gen0CollectionsEnd { get; private set; }
        public int Gen1CollectionsEnd { get; private set; }
        public int Gen2CollectionsEnd { get; private set; }
        public int Gen0CollectionsDelta { get; private set; }
        public int Gen1CollectionsDelta { get; private set; }
        public int Gen2CollectionsDelta { get; private set; }
        public bool OverBudget { get; private set; }
        public double BudgetMilliseconds { get; private set; }

        internal void AddRootNode(ScopeNode rootNode)
        {
            lock (rootNodesLock)
            {
                rootNodes.Add(rootNode);
            }
        }

        internal void End(double budgetMilliseconds)
        {
            EndTime = Stopwatch.GetTimestamp();
            BudgetMilliseconds = budgetMilliseconds;
            OverBudget = budgetMilliseconds > 0 && ElapsedMilliseconds() > budgetMilliseconds;

            AllocatedBytesEnd = GC.GetTotalAllocatedBytes(false);
            AllocatedBytesDelta = AllocatedBytesEnd - AllocatedBytesStart;

            Gen0CollectionsEnd = GC.CollectionCount(0);
            Gen1CollectionsEnd = GC.CollectionCount(1);
            Gen2CollectionsEnd = GC.CollectionCount(2);
            Gen0CollectionsDelta = Gen0CollectionsEnd - Gen0CollectionsStart;
            Gen1CollectionsDelta = Gen1CollectionsEnd - Gen1CollectionsStart;
            Gen2CollectionsDelta = Gen2CollectionsEnd - Gen2CollectionsStart;

            BuildLabelMetrics();
        }

        public double ElapsedMilliseconds()
        {
            long elapsedTicks = EndTime - StartTime;
            return ((double)(elapsedTicks * 1000)) / Stopwatch.Frequency;
        }

        private void BuildLabelMetrics()
        {
            labelMetrics.Clear();
            lock (rootNodesLock)
            {
                for (int i = 0; i < rootNodes.Count; i++)
                {
                    AccumulateLabelMetrics(rootNodes[i]);
                }
            }
        }

        private void AccumulateLabelMetrics(ScopeNode node)
        {
            double inclusiveMs = node.ElapsedMilliseconds();
            double selfMs = node.SelfMilliseconds();

            if (labelMetrics.TryGetValue(node.Label, out LabelMetrics current))
            {
                labelMetrics[node.Label] = new LabelMetrics(
                    current.InclusiveMs + inclusiveMs,
                    current.SelfMs + selfMs,
                    current.Calls + 1,
                    Math.Min(current.MinInclusiveMs, inclusiveMs),
                    Math.Max(current.MaxInclusiveMs, inclusiveMs));
            }
            else
            {
                labelMetrics[node.Label] = new LabelMetrics(inclusiveMs, selfMs, 1, inclusiveMs, inclusiveMs);
            }

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

    public class ScopeNode
    {
        public string Label { get; private set; } = string.Empty;
        public long StartTime { get; private set; }
        public long EndTime { get; private set; }
        public int ManagedThreadId { get; private set; }
        public List<ScopeNode> Children { get; } = new List<ScopeNode>();
        internal ScopeNode Parent { get; private set; }
        internal long ChildrenDurationTicks { get; private set; }

        internal void Reset(string label, int managedThreadId, ScopeNode parent)
        {
            Label = label;
            ManagedThreadId = managedThreadId;
            Parent = parent;
            StartTime = Stopwatch.GetTimestamp();
            EndTime = 0;
            ChildrenDurationTicks = 0;
            Children.Clear();
        }

        internal void End()
        {
            if (EndTime != 0)
            {
                return;
            }

            EndTime = Stopwatch.GetTimestamp();
            if (Parent != null)
            {
                Parent.ChildrenDurationTicks += Math.Max(0, EndTime - StartTime);
            }
        }

        public double ElapsedMilliseconds()
        {
            return ((double)(Math.Max(0, EndTime - StartTime))) * 1000 / Stopwatch.Frequency;
        }

        public double SelfMilliseconds()
        {
            long elapsedTicks = Math.Max(0, EndTime - StartTime);
            long selfTicks = Math.Max(0, elapsedTicks - ChildrenDurationTicks);
            return ((double)selfTicks) * 1000 / Stopwatch.Frequency;
        }

        // Add a child node (used for nested scopes)
        internal ScopeNode AddChild(string label)
        {
            ScopeNode child = RentNode(label, ManagedThreadId, this);
            Children.Add(child);
            return child;
        }
    }

    private const int maxFrames = 128;
    private const int rollingWindowSize = 240;

    public static bool IsRecording { get; private set; } = true;
    public static double FrameBudgetMilliseconds { get; set; } = 16.667;

    // Store only the last x amount of frames in memory.
    public static readonly BoundedQueue<Frame> Frames = new(maxFrames);

    // Use ThreadLocal to store a stack of ScopeNodes per thread and enable tracking of thread-local values.
    private static readonly ThreadLocal<Stack<ScopeNode>> threadLocalScopes = new ThreadLocal<Stack<ScopeNode>>(() => new Stack<ScopeNode>(), true);
    private static readonly ConcurrentDictionary<int, string> threadRootLabelCache = new ConcurrentDictionary<int, string>();
    private static readonly ConcurrentBag<ScopeNode> nodePool = new ConcurrentBag<ScopeNode>();
    private static readonly Dictionary<string, RollingWindow> rollingWindows = new Dictionary<string, RollingWindow>(256, StringComparer.Ordinal);
    private static readonly object rollingWindowsLock = new object();

    private static Frame currentFrame = null;
    private static uint frameCount = 0;

    public static void SetActive(bool isRecording)
    {
        if (IsRecording && !isRecording)
        {
            FinalizeCurrentFrame();
        }

        IsRecording = isRecording;
    }

    public static int CopyFramesTo(List<Frame> destination)
    {
        destination.Clear();
        foreach (Frame frame in Frames)
        {
            destination.Add(frame);
        }

        return destination.Count;
    }

    public static IReadOnlyDictionary<string, RollingLabelMetrics> GetRollingLabelMetricsSnapshot()
    {
        lock (rollingWindowsLock)
        {
            Dictionary<string, RollingLabelMetrics> snapshot = new Dictionary<string, RollingLabelMetrics>(rollingWindows.Count, StringComparer.Ordinal);
            foreach (KeyValuePair<string, RollingWindow> pair in rollingWindows)
            {
                snapshot[pair.Key] = pair.Value.Snapshot();
            }

            return snapshot;
        }
    }

    [Conditional("PROFILING")]
    public static void BeginFrame()
    {
        if (!IsRecording)
        {
            return;
        }

        if (currentFrame != null)
        {
            FinalizeCurrentFrame();
        }

        currentFrame = new Frame(frameCount);
    }

    [Conditional("PROFILING")]
    public static void EndFrame()
    {
        if (!IsRecording)
        {
            return;
        }

        FinalizeCurrentFrame();
    }

    [Conditional("PROFILING")]
    public static void BeginSample(string label)
    {
        if (!IsRecording || currentFrame == null)
        {
            return;
        }

        Stack<ScopeNode> scopes = threadLocalScopes.Value;  // Get the stack for the current thread
        Frame frame = currentFrame;
        if (frame == null)
        {
            return;
        }

        if (scopes.Count == 0)
        {
            // First scope for this thread (new root for this thread)
            int threadId = Environment.CurrentManagedThreadId;
            ScopeNode rootScopeNode = RentNode(GetThreadRootLabel(threadId), threadId, null);
            scopes.Push(rootScopeNode);
            frame.AddRootNode(rootScopeNode);
        }

        // Create a new child under the current top of the stack
        ScopeNode newScope = scopes.Peek().AddChild(label);

        scopes.Push(newScope);  // Push new scope to the thread's stack
    }

    [Conditional("PROFILING")]
    public static void EndSample()
    {
        if (!IsRecording || currentFrame == null)
        {
            return;
        }

        Stack<ScopeNode> scopes = threadLocalScopes.Value;

        if (scopes.Count > 1)
        {
            ScopeNode currentScope = scopes.Pop();
            currentScope.End();
        }
    }

    private static string GetThreadRootLabel(int threadId)
    {
        return threadRootLabelCache.GetOrAdd(threadId, static id => $"Thread-{id}");
    }

    private static ScopeNode RentNode(string label, int managedThreadId, ScopeNode parent)
    {
        if (!nodePool.TryTake(out ScopeNode node))
        {
            node = new ScopeNode();
        }

        node.Reset(label, managedThreadId, parent);
        return node;
    }

    private static void ReturnNodeTree(ScopeNode node)
    {
        for (int i = 0; i < node.Children.Count; i++)
        {
            ReturnNodeTree(node.Children[i]);
        }

        node.Reset(string.Empty, 0, null);
        nodePool.Add(node);
    }

    private static void FinalizeCurrentFrame()
    {
        Frame frame = currentFrame;
        if (frame == null)
        {
            return;
        }

        foreach (Stack<ScopeNode> scopes in threadLocalScopes.Values)
        {
            while (scopes.Count > 0)
            {
                ScopeNode currentScope = scopes.Pop();
                currentScope.End();
            }

            scopes.Clear();
        }

        frame.End(FrameBudgetMilliseconds);

        if (Frames.Enqueue(frame, out Frame evictedFrame))
        {
            for (int i = 0; i < evictedFrame.RootNodes.Count; i++)
            {
                ReturnNodeTree(evictedFrame.RootNodes[i]);
            }
        }

        UpdateRollingWindows(frame);
        frameCount++;
        currentFrame = null;
    }

    private static void UpdateRollingWindows(Frame frame)
    {
        lock (rollingWindowsLock)
        {
            foreach (KeyValuePair<string, LabelMetrics> pair in frame.LabelMetrics)
            {
                if (!rollingWindows.TryGetValue(pair.Key, out RollingWindow window))
                {
                    window = new RollingWindow(rollingWindowSize);
                    rollingWindows.Add(pair.Key, window);
                }

                window.Add(pair.Value.InclusiveMs);
            }
        }
    }
}