/*
 * Copyright 2018 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.
 */

namespace CartographerSharp.Common.Threading;

/// <summary>
/// Task state enumeration.
/// </summary>
public enum TaskState
{
    New,
    Dispatched,
    DependenciesCompleted,
    Running,
    Completed
}

/// <summary>
/// Represents a task that can be executed by a thread pool.
/// Tasks can have dependencies on other tasks.
/// </summary>
public class Task : IDisposable
{
    private readonly Lock _lock = new();
    private Action? _workItem;
    private ThreadPoolInterface? _threadPoolToNotify;
    private TaskState _state = TaskState.New;
    private int _uncompletedDependencies;
    private readonly HashSet<Task> _dependentTasks = [];
    private bool _disposed = false;

    /// <summary>
    /// Gets the current state of the task.
    /// </summary>
    public TaskState GetState()
    {
        lock (_lock)
        {
            return _state;
        }
    }

    /// <summary>
    /// Sets the work item to execute. State must be 'New'.
    /// </summary>
    public void SetWorkItem(Action workItem)
    {
        ArgumentNullException.ThrowIfNull(workItem);

        lock (_lock)
        {
            if (_state != TaskState.New)
            {
                throw new InvalidOperationException($"Cannot set work item when state is {_state}");
            }
            _workItem = workItem;
        }
    }

    /// <summary>
    /// Adds a dependency on another task. State must be 'New'.
    /// 'dependency' may be null, in which case it is assumed completed.
    /// </summary>
    public void AddDependency(WeakReference<Task>? dependency)
    {
        Task? depTask = null;
        lock (_lock)
        {
            if (_state != TaskState.New)
            {
                throw new InvalidOperationException($"Cannot add dependency when state is {_state}");
            }

            if (dependency != null && dependency.TryGetTarget(out depTask))
            {
                _uncompletedDependencies++;
            }
        }
        
        // Call AddDependentTask outside the lock to match C++ implementation
        depTask?.AddDependentTask(this);
    }

    /// <summary>
    /// Adds a dependent task that depends on this task.
    /// </summary>
    internal void AddDependentTask(Task dependentTask)
    {
        bool shouldNotifyImmediately = false;
        lock (_lock)
        {
            // If this task is already completed, notify the dependent task immediately
            if (_state == TaskState.Completed)
            {
                shouldNotifyImmediately = true;
            }
            else
            {
                if (!_dependentTasks.Add(dependentTask))
                {
                    throw new InvalidOperationException("Given dependency is already a dependency.");
                }
            }
        }
        
        // Call OnDependencyCompleted outside the lock to avoid potential deadlock
        if (shouldNotifyImmediately)
        {
            dependentTask.OnDependencyCompleted();
        }
    }

    /// <summary>
    /// Sets the thread pool that will execute this task. State must be 'New'.
    /// </summary>
    internal void SetThreadPool(ThreadPoolInterface threadPool)
    {
        lock (_lock)
        {
            if (_state != TaskState.New)
            {
                throw new InvalidOperationException($"Cannot set thread pool when state is {_state}");
            }

            _threadPoolToNotify = threadPool;
            _state = TaskState.Dispatched;
            
            if (_uncompletedDependencies == 0)
            {
                _state = TaskState.DependenciesCompleted;
                _threadPoolToNotify.NotifyDependenciesCompleted(this);
            }
        }
    }

    /// <summary>
    /// Called when a dependency completes. State must be 'New' or 'Dispatched'.
    /// If 'Dispatched', may become 'DependenciesCompleted'.
    /// </summary>
    internal void OnDependencyCompleted()
    {
        ThreadPoolInterface? threadPoolToNotify = null;
        lock (_lock)
        {
            if (_state != TaskState.New && _state != TaskState.Dispatched)
            {
                throw new InvalidOperationException($"Cannot complete dependency when state is {_state}");
            }

            _uncompletedDependencies--;
            if (_uncompletedDependencies == 0 && _state == TaskState.Dispatched)
            {
                _state = TaskState.DependenciesCompleted;
                threadPoolToNotify = _threadPoolToNotify;
            }
        }
        
        threadPoolToNotify?.NotifyDependenciesCompleted(this);
    }

    /// <summary>
    /// Executes the task. State must be 'DependenciesCompleted' and becomes 'Completed'.
    /// </summary>
    internal void Execute()
    {
        Action? workItem;
        lock (_lock)
        {
            if (_state != TaskState.DependenciesCompleted)
            {
                throw new InvalidOperationException($"Cannot execute task when state is {_state}");
            }
            _state = TaskState.Running;
            workItem = _workItem;
        }

        // Execute the work item outside the lock
        try
        {
            workItem?.Invoke();
        }
        finally
        {
            // Set state to Completed and notify all dependent tasks in a single lock
            lock (_lock)
            {
                _state = TaskState.Completed;
                
                // Notify all dependent tasks
                foreach (var dependentTask in _dependentTasks)
                {
                    dependentTask.OnDependencyCompleted();
                }
            }
        }
    }

    public void Dispose()
    {
        if (!_disposed)
        {
            _disposed = true;
        }

        GC.SuppressFinalize(this);
    }
}