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 class Frame(uint frameCount)
    {
        public uint FrameCount { get; } = frameCount;
        public long StartTime { get; } = Stopwatch.GetTimestamp();
        public long EndTime { get; private set; }

        // Use a concurrent list to collect all thread root nodes per frame.
        public ConcurrentBag<ScopeNode> RootNodes = new ConcurrentBag<ScopeNode>();

        internal void End()
        {
            EndTime = Stopwatch.GetTimestamp();
        }

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

    public class ScopeNode(string label)
    {
        public string Label { get; } = label;
        public long StartTime { get; private set; } = Stopwatch.GetTimestamp();  // Start time in ticks
        public long EndTime { get; private set; }
        public int ManagedThreadId { get; } = Environment.CurrentManagedThreadId;
        public List<ScopeNode> Children { get; } = new List<ScopeNode>();

        internal void End()
        {
            EndTime = Stopwatch.GetTimestamp();  // End time in ticks
        }

        public double ElapsedMilliseconds()
        {
            return ((double)(EndTime - StartTime)) * 1000 / Stopwatch.Frequency;  // Convert ticks to ms
        }

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

    private const int maxFrames = 128;

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

    // 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 Frame currentFrame = null;
    private static uint frameCount = 0;

    public static void SetActive(bool isRecording)
    {
        IsRecording = isRecording;
    }

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

        currentFrame = new Frame(frameCount);
    }

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

        foreach (Stack<ScopeNode> scopes in threadLocalScopes.Values)
        {
            if (scopes.Count > 0)
            {
                // Pop the left over root nodes.
                ScopeNode currentScope = scopes.Pop();
                currentScope.End();
            }

            // Clean up the thread-local stack to ensure it's empty for the next frame.
            scopes.Clear();
        }

        currentFrame.End();
        Frames.Enqueue(currentFrame);
        frameCount++;
    }

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

        Stack<ScopeNode> scopes = threadLocalScopes.Value;  // Get the stack for the current thread

        if (scopes.Count == 0)
        {
            // First scope for this thread (new root for this thread)
            ScopeNode rootScopeNode = new ScopeNode($"Thread-{Environment.CurrentManagedThreadId}");
            scopes.Push(rootScopeNode);
            currentFrame.RootNodes.Add(rootScopeNode);  // Add root node to the frame list
        }

        // 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)
        {
            return;
        }

        Stack<ScopeNode> scopes = threadLocalScopes.Value;

        if (scopes.Count > 0)
        {
            // Only pop if this is not the root node.
            //ScopeNode currentScope = scopes.Count > 1 ? scopes.Pop() : scopes.Peek();
            ScopeNode currentScope = scopes.Pop();
            currentScope.End();
        }
    }
}