ObservableTask.cs

using CommunityToolkit.Mvvm.ComponentModel;
using CommunityToolkit.Mvvm.Input;
using NeuroAccessMaui.Services; // For LogService
using NeuroAccessMaui.Services.Resilience.Dispatch;
using System.Threading.Tasks;
using System.Threading; // Added for OperationCanceledException
 
namespace NeuroAccessMaui.UI.MVVM
{
    public enum ObservableTaskStatus
    {
        Pending,
        Running,
        Succeeded,
        Failed,
        Canceled
    }
 
    public partial class ObservableTask<TProgress> : ObservableObject, IDisposable
    {
        private readonly object syncObject = new();
        private readonly List<IDisposable> attachedDisposables = new();
 
        // Backing field for the latest progress value.
        private TProgress progress = default!;
 
        // Commands to notify when state changes.
        private readonly List<IRelayCommand> notifyCommands = new();
 
        // The factory delegate to create tasks; stored so that a new load can be triggered.
        private Func<TaskContext<TProgress>, Task>? taskFactory;
 
        // A generation counter to identify the currently active task.
        private int currentGeneration = 0;
 
        public bool UseTaskRun { get; set; } = false;
 
        public Task? CurrentTask { get; private set; }
 
        public Task? CurrentWatcher { get; private set; }
 
        public IDispatcherAdapter Dispatcher { get; set; } = UiDispatcher.Instance;
 
        public ObservableTaskStatus State { get; private set; }
 
        // Convenience properties for data binding.
        public bool IsPending => this.State == ObservableTaskStatus.Pending;
        public bool IsNotPending => this.State != ObservableTaskStatus.Pending;
        public bool IsRunning => this.State == ObservableTaskStatus.Running;
        public bool IsNotRunning => this.State != ObservableTaskStatus.Running;
        public bool IsSucceeded => this.State == ObservableTaskStatus.Succeeded;
        public bool IsNotSucceeded => this.State != ObservableTaskStatus.Succeeded;
        public bool IsFailed => this.State == ObservableTaskStatus.Failed;
        public bool IsNotFailed => this.State != ObservableTaskStatus.Failed;
        public bool IsCanceled => this.State == ObservableTaskStatus.Canceled;
        public bool IsNotCanceled => this.State != ObservableTaskStatus.Canceled;
        public bool IsRefreshing { get; private set; }
 
        public AggregateException? Exception { get; private set; }
 
        public Exception? InnerException { get; private set; }
 
        public string? ErrorMessage { get; private set; }
 
        public TProgress Progress
        {
            get => this.progress;
            private set
            {
                if (!EqualityComparer<TProgress>.Default.Equals(this.progress, value))
                {
                    this.progress = value;
                    // Marshal via dispatcher.
                this.Dispatcher.Post(() => this.OnPropertyChanged(nameof(this.Progress)));
                }
            }
        }
 
        public CancellationTokenSource? CancellationTokenSource { get; private set; }
 
        public event EventHandler<ObservableTaskStatus>? StateChanged;
 
        private void CancelExistingTask()
        {
            CancellationTokenSource? Cts = this.CancellationTokenSource;
            if (Cts is null)
                return;
 
            try
            {
                if (this.CurrentTask is { IsCompleted: false })
                {
                    Cts.Cancel();
                }
            }
            catch (Exception ex)
            {
                ServiceRef.LogService.LogException(ex);
            }
        }
 
        private void StartNewTask(bool isRefresh)
        {
            try
            {
                lock (this.syncObject)
                {
                    if (this.taskFactory is null)
                        return;
 
                    // Increment the generation token.
                    int ThisGeneration = ++this.currentGeneration;
 
                    // Cancel any existing task.
                    this.CancelExistingTask();
 
                    this.CancellationTokenSource = new CancellationTokenSource();
                    this.CurrentTask = null;
                    this.CurrentWatcher = null;
                    this.ErrorMessage = string.Empty;
                    this.Exception = null;
                    this.InnerException = null;
                    this.IsRefreshing = isRefresh;
                    this.State = ObservableTaskStatus.Running;
 
                    // Create a progress reporter that ensures UI updates and ignores stale generations.
                    Progress<TProgress> ProgressReporter = new(value =>
                    {
                        if (ThisGeneration == this.currentGeneration)
                        {
                            this.Dispatcher.Post(() => this.Progress = value);
                        }
                    });
 
                    TaskContext<TProgress> Context = new(isRefresh, this.CancellationTokenSource.Token, ProgressReporter);
 
                    // Start the new task.
                    if (this.UseTaskRun)
                    {
                        // Force background-thread execution (good for CPU-bound work).
                        this.CurrentTask = Task.Run(async () => await this.taskFactory!(Context), this.CancellationTokenSource.Token);
                    }
                    else
                    {
                        // Default path (good for naturally async/IO-bound work).
                        this.CurrentTask = this.taskFactory.Invoke(Context);
                    }
 
                    this.CurrentWatcher = this.WatchTaskAsync(ThisGeneration, this.CurrentTask, this.CancellationTokenSource);
                }
            }
            catch (Exception ex)
            {
                ServiceRef.LogService.LogException(ex);
            }
            finally
            {
                // Notify bindings and start watching the new task.
                this.NotifyAll();
            }
        }
 
        private async Task WatchTaskAsync(int generation, Task? current, CancellationTokenSource cts)
        {
            // TODO: Add logging for tasks gone rogue if desired.
 
            AggregateException? AggregateException = null;
            Exception? InnerException = null;
            string? ErrorMessage = null;
 
            try
            {
                if (current is not null)
                {
                    await current.ConfigureAwait(false);
                }
            }
            catch (OperationCanceledException)
            {
                // Normal cooperative cancellation (includes TaskCanceledException). Suppress error logging.
            }
            catch (Exception ex)
            {
                // Exceptions are surfaced via Exception/InnerException/ErrorMessage.
                AggregateException = ex as AggregateException;
                InnerException = ex.InnerException ?? ex;
                ErrorMessage = InnerException.Message;
 
                // Log by default (can be routed to telemetry later).
                ServiceRef.LogService.LogException(ex);
            }
            finally
            {
                // Marshal via dispatcher for state updates.
                this.Dispatcher.Post(() =>
                {
                    lock (this.syncObject)
                    {
                        cts.Dispose();
 
                        if (generation != this.currentGeneration)
                            return;
 
                        if (current is null)
                            this.State = ObservableTaskStatus.Pending;
                        else if (!current.IsCompleted)
                            this.State = ObservableTaskStatus.Running;
                        else if (current.IsCanceled)
                            this.State = ObservableTaskStatus.Canceled;
                        else if (current.IsFaulted)
                            this.State = ObservableTaskStatus.Failed;
                        else
                            this.State = ObservableTaskStatus.Succeeded;
 
                        this.Exception = AggregateException;
                        this.InnerException = InnerException;
                        this.ErrorMessage = ErrorMessage;
 
                        this.IsRefreshing = false;
 
                        if (this.CancellationTokenSource is not null)
                        {
                            this.CancellationTokenSource = null;
                        }
                    }
 
                    this.NotifyAll();
                    this.StateChanged?.Invoke(this, this.State);
                });
            }
        }
 
        private void NotifyAll()
        {
            this.Dispatcher.Post(() =>
            {
                this.OnPropertyChanged(nameof(this.State));
                this.OnPropertyChanged(nameof(this.IsPending));
                this.OnPropertyChanged(nameof(this.IsNotPending));
                this.OnPropertyChanged(nameof(this.IsRunning));
                this.OnPropertyChanged(nameof(this.IsNotRunning));
                this.OnPropertyChanged(nameof(this.IsSucceeded));
                this.OnPropertyChanged(nameof(this.IsNotSucceeded));
                this.OnPropertyChanged(nameof(this.IsFailed));
                this.OnPropertyChanged(nameof(this.IsNotFailed));
                this.OnPropertyChanged(nameof(this.IsCanceled));
                this.OnPropertyChanged(nameof(this.IsNotCanceled));
                this.OnPropertyChanged(nameof(this.Exception));
                this.OnPropertyChanged(nameof(this.InnerException));
                this.OnPropertyChanged(nameof(this.ErrorMessage));
                this.OnPropertyChanged(nameof(this.Progress));
                this.OnPropertyChanged(nameof(this.IsRefreshing));
 
                foreach (IRelayCommand Command in this.notifyCommands)
                {
                    Command.NotifyCanExecuteChanged();
                }
 
                // Also update the automatically generated commands.
                this.CancelCommand.NotifyCanExecuteChanged();
                this.ReloadCommand.NotifyCanExecuteChanged();
                this.RefreshCommand.NotifyCanExecuteChanged();
            });
        }
 
        public async Task WaitCurrentAsync()
        {
            Task? Watcher = this.CurrentWatcher;
            if (Watcher is null)
                return;
 
            await Watcher.ConfigureAwait(false);
        }
 
        public async Task WaitAllAsync(bool waitWhilePending = false)
        {
            while (true)
            {
                Task? Watcher = this.CurrentWatcher;
                if (Watcher is null)
                {
                    if (waitWhilePending && this.State == ObservableTaskStatus.Pending)
                    {
                        await Task.Delay(Constants.Delays.Default).ConfigureAwait(false);
                        continue;
                    }
 
                    return;
                }
 
                await Watcher.ConfigureAwait(false);
 
                Task? NewWatcher = this.CurrentWatcher;
                if (NewWatcher is not null && !ReferenceEquals(Watcher, NewWatcher))
                    continue;
 
                break;
            }
        }
 
        public void Configure(Func<TaskContext<TProgress>, Task> task, params IRelayCommand[] notifyCommands)
        {
            lock (this.syncObject)
            {
                this.taskFactory = task;
                this.notifyCommands.Clear();
                if (notifyCommands.Length > 0)
                    this.notifyCommands.AddRange(notifyCommands);
            }
 
            // Leave state as-is (likely Pending). No auto-start here.
            this.NotifyAll();
        }
 
        public void Configure<TArg>(Func<TArg, TaskContext<TProgress>, Task> task, TArg arg, params IRelayCommand[] notifyCommands)
            => this.Configure(ctx => task(arg, ctx), notifyCommands);
 
        public void Load(Func<TaskContext<TProgress>, Task> task, params IRelayCommand[] notifyCommands)
        {
            this.Configure(task, notifyCommands);
            this.StartNewTask(isRefresh: false);
        }
 
        public void Load<TArg>(Func<TArg, TaskContext<TProgress>, Task> task, TArg arg, params IRelayCommand[] notifyCommands)
        {
            this.Configure(ctx => task(arg, ctx), notifyCommands);
            this.StartNewTask(isRefresh: false);
        }
 
        public void Load(Func<Task> task, params IRelayCommand[] notifyCommands)
        {
            // Wrap the parameterless task factory in a lambda that ignores the TaskContext.
            this.Load((Func<TaskContext<TProgress>, Task>)(_ => task()), notifyCommands);
        }
 
        public void Load<TArg>(Func<TArg, Task> task, TArg arg, params IRelayCommand[] notifyCommands)
        {
            this.Load((Func<TaskContext<TProgress>, Task>)(_ => task(arg)), notifyCommands);
        }
 
        public void Run() => this.StartNewTask(isRefresh: false);
 
        public void RunRefresh() => this.StartNewTask(isRefresh: true);
 
        public void Run<TArg>(Func<TArg, TaskContext<TProgress>, Task> task, TArg arg)
        {
            lock (this.syncObject)
            {
                this.taskFactory = ctx => task(arg, ctx);
            }
            this.StartNewTask(isRefresh: false);
        }
 
        public void RunRefresh<TArg>(Func<TArg, TaskContext<TProgress>, Task> task, TArg arg)
        {
            lock (this.syncObject)
            {
                this.taskFactory = ctx => task(arg, ctx);
            }
            this.StartNewTask(isRefresh: true);
        }
 
        [RelayCommand(CanExecute = nameof(IsNotPending))]
        public void Reload()
        {
            if (this.IsPending)
                return;
 
            this.StartNewTask(isRefresh: false);
        }
 
        [RelayCommand(CanExecute = nameof(IsNotPending))]
        public void Refresh()
        {
            if (this.IsPending)
                return;
 
            this.StartNewTask(isRefresh: true);
        }
 
        [RelayCommand(CanExecute = nameof(IsRunning))]
        public void Cancel()
        {
            if (!this.IsRunning)
                return;
 
            lock (this.syncObject)
            {
                try
                {
                    this.CancellationTokenSource?.Cancel();
                }
                catch
                {
                    // Ignored.
                }
            }
        }
 
        private bool disposed;
        // Public Dispose: callers clean up managed + unmanaged
        public void Dispose()
        {
            this.Dispose(disposing: true);
            GC.SuppressFinalize(this);
        }
 
        // Finalizer only if you manage native resources.
        ~ObservableTask()
        {
            this.Dispose(disposing: false);
        }
 
        // The core pattern
        protected virtual void Dispose(bool disposing)
        {
            if (this.disposed)
                return;
            this.disposed = true;
 
            if (disposing)
            {
                lock (this.syncObject)
                {
                    try { this.CancellationTokenSource?.Cancel(); }
                    catch { /* ignore */ }
 
                    this.CancellationTokenSource?.Dispose();
                    this.CancellationTokenSource = null;
 
                    // Dispose any attached disposables (e.g., policies)
                    foreach (var d in this.attachedDisposables)
                    {
                        try { d.Dispose(); }
                        catch { /* ignore */ }
                    }
                    this.attachedDisposables.Clear();
                }
            }
        }
 
        internal void AttachDisposable(IDisposable disposable)
        {
            if (disposable is null)
                return;
            lock (this.syncObject)
            {
                this.attachedDisposables.Add(disposable);
            }
        }
 
    }
 
    public partial class ObservableTask : ObservableTask<int>
    {
    }
 
    public sealed class ObservableTask<TResult, TProgress> : ObservableTask<TProgress>
    {
        public TResult? Result { get; private set; }
 
        public void Load(Func<TaskContext<TProgress>, Task<TResult>> op, params IRelayCommand[] notify)
        {
            base.Load(async ctx =>
            {
                this.Result = await op(ctx);
            }, notify);
        }
 
        public void Configure(Func<TaskContext<TProgress>, Task<TResult>> op, params IRelayCommand[] notify)
            => base.Configure(async ctx => { this.Result = await op(ctx); }, notify);
 
        public void Configure<TArg>(Func<TArg, TaskContext<TProgress>, Task<TResult>> op, TArg arg, params IRelayCommand[] notify)
            => base.Configure(async ctx => { this.Result = await op(arg, ctx); }, notify);
 
        public void Load<TArg>(Func<TArg, TaskContext<TProgress>, Task<TResult>> op, TArg arg, params IRelayCommand[] notify)
        {
            base.Load(async ctx => { this.Result = await op(arg, ctx); }, notify);
        }
 
        public void Run<TArg>(Func<TArg, TaskContext<TProgress>, Task<TResult>> op, TArg arg)
        {
            base.Run<TArg>(async (a, ctx) => { this.Result = await op(a, ctx); }, arg);
        }
 
        public void RunRefresh<TArg>(Func<TArg, TaskContext<TProgress>, Task<TResult>> op, TArg arg)
        {
            base.RunRefresh<TArg>(async (a, ctx) => { this.Result = await op(a, ctx); }, arg);
        }
    }
 
}