using System.Diagnostics;

namespace Nerfed.Runtime;

public sealed class JobSystem : IDisposable
{
    public static readonly JobSystem Default = new JobSystem();

    private Thread[] workers;
    private SemaphoreSlim startSignal;
    private CountdownEvent completionEvent;
    private volatile bool running;

    // Shared per-dispatch state written by main thread before workers wake.
    private volatile Action<int> currentAction;
    private int workCount;
    private int nextIndex; // grabbed with Interlocked.Increment for work-stealing

    public int WorkerCount => workers?.Length ?? 0;
    public bool IsInitialized => workers != null;

    public void Initialize(int threadCount = -1)
    {
        if (IsInitialized)
        {
            throw new InvalidOperationException("JobSystem is already initialized. Call Shutdown first.");
        }

        threadCount = threadCount < 0
            ? Math.Max(1, Environment.ProcessorCount - 2)
            : Math.Max(1, threadCount);

        running = true;
        startSignal = new SemaphoreSlim(0);
        completionEvent = new CountdownEvent(1);
        workers = new Thread[threadCount];

        for (int i = 0; i < threadCount; i++)
        {
            workers[i] = new Thread(WorkerLoop)
            {
                IsBackground = true,
                Name = $"Job-{i}",
            };
            workers[i].Start();
        }
    }

    public void Dispatch(int count, Action<int> action)
    {
        if (count <= 0) return;

        if (!IsInitialized)
        {
            // Safe fallback: run inline if Initialize was never called.
            for (int i = 0; i < count; i++)
                action(i);
            return;
        }

        currentAction = action;
        workCount = count;
        Volatile.Write(ref nextIndex, 0);

        completionEvent.Reset(workers.Length);
        startSignal.Release(workers.Length);
        completionEvent.Wait();

        currentAction = null;
    }

    public void Shutdown()
    {
        if (!IsInitialized) return;

        running = false;
        startSignal.Release(workers.Length); // wake all workers so they can see running=false and exit

        foreach (Thread t in workers)
            t.Join();

        completionEvent.Dispose();
        startSignal.Dispose();
        workers = null;
    }

    public void Dispose() => Shutdown();

    private void WorkerLoop()
    {
        while (true)
        {
            startSignal.Wait();
            if (!running) return;

            Action<int> action = currentAction;
            int total = workCount;

            while (true)
            {
                int index = Interlocked.Increment(ref nextIndex) - 1;
                if (index >= total) break;
                action(index);
            }

            completionEvent.Signal();
        }
    }
}