StabilityMatrix/StabilityMatrix.Core/Models/TrackedDownload.cs

using System.Diagnostics.CodeAnalysis;
using System.Text.Json.Serialization;
using AsyncAwaitBestPractices;
using NLog;
using StabilityMatrix.Core.Helper;
using StabilityMatrix.Core.Models.FileInterfaces;
using StabilityMatrix.Core.Models.Progress;
using StabilityMatrix.Core.Services;

namespace StabilityMatrix.Core.Models;

[JsonSerializable(typeof(TrackedDownload))]
public class TrackedDownload
{
    private static readonly Logger Logger = LogManager.GetCurrentClassLogger();
    
    [JsonIgnore]
    private IDownloadService? downloadService;
    
    [JsonIgnore]
    private Task? downloadTask;
    
    [JsonIgnore]
    private CancellationTokenSource? downloadCancellationTokenSource;
    
    [JsonIgnore]
    private CancellationTokenSource? downloadPauseTokenSource;
    
    [JsonIgnore]
    private CancellationTokenSource AggregateCancellationTokenSource =>
        CancellationTokenSource.CreateLinkedTokenSource(
            downloadCancellationTokenSource?.Token ?? CancellationToken.None,
            downloadPauseTokenSource?.Token ?? CancellationToken.None);
    
    public required Guid Id { get; init; }
    
    public required Uri SourceUrl { get; init; }
    
    public Uri? RedirectedUrl { get; init; }
    
    public required DirectoryPath DownloadDirectory { get; init; }
    
    public required string FileName { get; init; }
    
    public required string TempFileName { get; init; }
    
    public string? ExpectedHashSha256 { get; set; }
    
    [JsonIgnore]
    [MemberNotNullWhen(true, nameof(ExpectedHashSha256))]
    public bool ValidateHash => ExpectedHashSha256 is not null;

    [JsonConverter(typeof(JsonStringEnumConverter))]
    public ProgressState ProgressState { get; set; } = ProgressState.Inactive;

    public List<string> ExtraCleanupFileNames { get; init; } = new();
    
    // Used for restoring progress on load
    public long DownloadedBytes { get; set; }
    public long TotalBytes { get; set; }
    
    /// <summary>
    /// Optional context action to be invoked on completion
    /// </summary>
    public IContextAction? ContextAction { get; set; }
    
    [JsonIgnore]
    public Exception? Exception { get; private set; }
    
    #region Events
    private WeakEventManager<ProgressReport>? progressUpdateEventManager;
    
    public event EventHandler<ProgressReport> ProgressUpdate
    {
        add
        {
            progressUpdateEventManager ??= new WeakEventManager<ProgressReport>();
            progressUpdateEventManager.AddEventHandler(value);
        }
        remove => progressUpdateEventManager?.RemoveEventHandler(value);
    }
    
    protected void OnProgressUpdate(ProgressReport e)
    {
        // Update downloaded and total bytes
        DownloadedBytes = Convert.ToInt64(e.Current);
        TotalBytes = Convert.ToInt64(e.Total);
        
        progressUpdateEventManager?.RaiseEvent(this, e, nameof(ProgressUpdate));
    }
    
    private WeakEventManager<ProgressState>? progressStateChangedEventManager;
    
    public event EventHandler<ProgressState> ProgressStateChanged
    {
        add
        {
            progressStateChangedEventManager ??= new WeakEventManager<ProgressState>();
            progressStateChangedEventManager.AddEventHandler(value);
        }
        remove => progressStateChangedEventManager?.RemoveEventHandler(value);
    }
    
    protected void OnProgressStateChanged(ProgressState e)
    {
        progressStateChangedEventManager?.RaiseEvent(this, e, nameof(ProgressStateChanged));
    }
    #endregion
    
    [MemberNotNull(nameof(downloadService))]
    private void EnsureDownloadService()
    {
        if (downloadService == null)
        {
            throw new InvalidOperationException("Download service is not set.");
        }
    }

    private async Task StartDownloadTask(long resumeFromByte, CancellationToken cancellationToken)
    {
        var progress = new Progress<ProgressReport>(OnProgressUpdate);
        
        await downloadService!.ResumeDownloadToFileAsync(
            SourceUrl.ToString(),
            DownloadDirectory.JoinFile(TempFileName),
            resumeFromByte,
            progress,
            cancellationToken: cancellationToken).ConfigureAwait(false);
        
        // If hash validation is enabled, validate the hash
        if (ValidateHash)
        {
            OnProgressUpdate(new ProgressReport(0, isIndeterminate: true, type: ProgressType.Hashing));
            var hash = await FileHash.GetSha256Async(DownloadDirectory.JoinFile(TempFileName), progress).ConfigureAwait(false);
            if (hash != ExpectedHashSha256?.ToLowerInvariant())
            {
                throw new Exception($"Hash validation for {FileName} failed, expected {ExpectedHashSha256} but got {hash}");
            }
        }
    }
    
    public void Start()
    {
        if (ProgressState != ProgressState.Inactive)
        {
            throw new InvalidOperationException($"Download state must be inactive to start, not {ProgressState}");
        }
        Logger.Debug("Starting download {Download}", FileName);
        
        EnsureDownloadService();
        
        downloadCancellationTokenSource = new CancellationTokenSource();
        downloadPauseTokenSource = new CancellationTokenSource();
        
        downloadTask = StartDownloadTask(0, AggregateCancellationTokenSource.Token)
            .ContinueWith(OnDownloadTaskCompleted);
        
        ProgressState = ProgressState.Working;
        OnProgressStateChanged(ProgressState);
    }

    public void Resume()
    {
        if (ProgressState != ProgressState.Inactive)
        {
            Logger.Warn("Attempted to resume download {Download} but it is not paused ({State})", FileName, ProgressState);
        }
        Logger.Debug("Resuming download {Download}", FileName);
        
        // Read the temp file to get the current size
        var tempSize = 0L;
        
        var tempFile = DownloadDirectory.JoinFile(TempFileName);
        if (tempFile.Exists)
        {
            tempSize = tempFile.Info.Length;
        }
        
        EnsureDownloadService();
        
        downloadCancellationTokenSource = new CancellationTokenSource();
        downloadPauseTokenSource = new CancellationTokenSource();
        
        downloadTask = StartDownloadTask(tempSize, AggregateCancellationTokenSource.Token)
            .ContinueWith(OnDownloadTaskCompleted);
        
        ProgressState = ProgressState.Working;
        OnProgressStateChanged(ProgressState);
    }

    public void Pause()
    {
        if (ProgressState != ProgressState.Working)
        {
            Logger.Warn("Attempted to pause download {Download} but it is not in progress ({State})", FileName, ProgressState);
            return;
        }
        
        Logger.Debug("Pausing download {Download}", FileName);
        downloadPauseTokenSource?.Cancel();
    }
    
    public void Cancel()
    {
        if (ProgressState is not (ProgressState.Working or ProgressState.Inactive))
        {
            Logger.Warn("Attempted to cancel download {Download} but it is not in progress ({State})", FileName, ProgressState);
            return;
        }
        
        Logger.Debug("Cancelling download {Download}", FileName);

        // Cancel token if it exists
        if (downloadCancellationTokenSource is { } token)
        {
            token.Cancel();
        }
        // Otherwise handle it manually
        else
        {
            DoCleanup();
            
            ProgressState = ProgressState.Cancelled;
            OnProgressStateChanged(ProgressState);
        }
    }

    /// <summary>
    /// Deletes the temp file and any extra cleanup files
    /// </summary>
    private void DoCleanup()
    {
        try
        {
            DownloadDirectory.JoinFile(TempFileName).Delete();
        }
        catch (IOException)
        {
            Logger.Warn("Failed to delete temp file {TempFile}", TempFileName);
        }
        
        foreach (var extraFile in ExtraCleanupFileNames)
        {
            try
            {
                DownloadDirectory.JoinFile(extraFile).Delete();
            }
            catch (IOException)
            {
                Logger.Warn("Failed to delete extra cleanup file {ExtraFile}", extraFile);
            }
        }
    }
    
    /// <summary>
    /// Invoked by the task's completion callback
    /// </summary>
    private void OnDownloadTaskCompleted(Task task)
    {
        // For cancelled, check if it was actually cancelled or paused
        if (task.IsCanceled)
        {
            // If the task was cancelled, set the state to cancelled
            if (downloadCancellationTokenSource?.IsCancellationRequested == true)
            {
                ProgressState = ProgressState.Cancelled;
            }
            // If the task was not cancelled, set the state to paused
            else if (downloadPauseTokenSource?.IsCancellationRequested == true)
            {
                ProgressState = ProgressState.Inactive;
            }
            else
            {
                throw new InvalidOperationException("Download task was cancelled but neither cancellation token was cancelled.");
            }
        }
        // For faulted
        else if (task.IsFaulted)
        {
            // Set the exception
            Exception = task.Exception;

            ProgressState = ProgressState.Failed;
        }
        // Otherwise success
        else
        {
            ProgressState = ProgressState.Success;
        }

        // For failed or cancelled, delete the temp files
        if (ProgressState is ProgressState.Failed or ProgressState.Cancelled)
        {
            DoCleanup();
        } 
        else if (ProgressState == ProgressState.Success)
        {
            // Move the temp file to the final file
            DownloadDirectory.JoinFile(TempFileName).MoveTo(DownloadDirectory.JoinFile(FileName));
        }
        
        // For pause, just do nothing
        
        OnProgressStateChanged(ProgressState);
        
        // Dispose of the task and cancellation token
        downloadTask = null;
        downloadCancellationTokenSource = null;
        downloadPauseTokenSource = null;
    }
    
    public void SetDownloadService(IDownloadService service)
    {
        downloadService = service;
    }
}