/*
 * Copyright 2016 The Cartographer Authors
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

using Microsoft.Extensions.Logging;
using System.Runtime.InteropServices;

namespace CartographerSharp.Common.Threading;

/// <summary>
/// Interface for thread pool implementations.
/// </summary>
public abstract class ThreadPoolInterface
{
    /// <summary>
    /// Schedules a task for execution.
    /// </summary>
    /// <param name="task">The task to schedule</param>
    /// <returns>A weak reference to the task</returns>
    public abstract WeakReference<Task> Schedule(Task task);

    /// <summary>
    /// Executes a task.
    /// </summary>
    protected static void Execute(Task task)
    {
        task.Execute();
    }

    /// <summary>
    /// Sets the thread pool for a task.
    /// </summary>
    protected void SetThreadPool(Task task)
    {
        task.SetThreadPool(this);
    }

    /// <summary>
    /// Notifies that dependencies of a task are completed.
    /// </summary>
    internal abstract void NotifyDependenciesCompleted(Task task);
}

/// <summary>
/// A fixed number of threads working on tasks. Adding a task does not block.
/// Tasks may be added whether or not their dependencies are completed.
/// When all dependencies of a task are completed, it is queued up for execution
/// in a background thread. The queue must be empty before calling the destructor.
/// The thread pool will then wait for the currently executing work items to finish
/// and then destroy the threads.
/// </summary>
public partial class ThreadPool : ThreadPoolInterface, IDisposable
{
    private readonly object _lock = new();
    private readonly List<Thread> _pool = [];
    private readonly Queue<Task> _taskQueue = new();
    private readonly Dictionary<Task, Task> _tasksNotReady = [];
    private bool _running = true;
    private bool _disposed = false;

    /// <summary>
    /// Creates a new thread pool with the specified number of threads.
    /// </summary>
    public ThreadPool(int numThreads)
    {
        if (numThreads <= 0)
        {
            throw new ArgumentException("ThreadPool requires a positive num_threads!", nameof(numThreads));
        }

        lock (_lock)
        {
            for (int i = 0; i < numThreads; i++)
            {
                var thread = new Thread(DoWork)
                {
                    IsBackground = false,
                    Name = $"CartographerThreadPool-{i}"
                };
                _pool.Add(thread);
                thread.Start();
            }
        }
    }

    /// <summary>
    /// Schedules a task for execution.
    /// </summary>
    public override WeakReference<Task> Schedule(Task task)
    {
        ArgumentNullException.ThrowIfNull(task);

        lock (_lock)
        {
            if (!_running)
            {
                throw new InvalidOperationException("ThreadPool is not running");
            }

            if (_tasksNotReady.ContainsKey(task))
            {
                throw new InvalidOperationException("Schedule called twice for the same task");
            }

            _tasksNotReady[task] = task;
        }

        SetThreadPool(task);
        return new WeakReference<Task>(task);
    }

    /// <summary>
    /// Notifies that dependencies of a task are completed.
    /// </summary>
    internal override void NotifyDependenciesCompleted(Task task)
    {
        lock (_lock)
        {
            if (!_tasksNotReady.TryGetValue(task, out var storedTask))
            {
                throw new InvalidOperationException("Task not found in tasks_not_ready");
            }

            _taskQueue.Enqueue(storedTask);
            _tasksNotReady.Remove(task);
            Monitor.PulseAll(_lock);
        }
    }

    /// <summary>
    /// Worker thread method.
    /// </summary>
    private void DoWork()
    {
        Thread.BeginThreadAffinity();
        try
        {
            while (true)
            {
                Task? task = null;

                lock (_lock)
                {
                    while (_taskQueue.Count == 0 && _running)
                    {
                        Monitor.Wait(_lock);
                    }

                    if (!_running && _taskQueue.Count == 0)
                    {
                        return;
                    }

                    if (_taskQueue.Count > 0)
                    {
                        task = _taskQueue.Dequeue();
                    }
                }

                if (task != null)
                {
                    if (task.GetState() != TaskState.DependenciesCompleted)
                    {
                        continue;
                    }

                    Execute(task);
                }
            }
        }
        finally
        {
            Thread.EndThreadAffinity();
        }
    }

    /// <summary>
    /// Disposes the thread pool, waiting for all tasks to complete.
    /// </summary>
    public void Dispose()
    {
        if (_disposed)
        {
            return;
        }

        lock (_lock)
        {
            if (!_running)
            {
                return;
            }

            _running = false;
            Monitor.PulseAll(_lock);
        }

        foreach (var thread in _pool)
        {
            thread.Join();
        }

        _disposed = true;

        GC.SuppressFinalize(this);
    }

    [LibraryImport("libc", SetLastError = true)]
    private static partial int nice(int inc);
}