MilkyShots/Lactose/Jobs/ThumbnailJob.cs

using Butter.Settings;
using Lactose.Models;
using Lactose.Repositories;
using SixLabors.ImageSharp;
using SixLabors.ImageSharp.Formats.Webp;
using SixLabors.ImageSharp.Processing;

namespace Lactose.Jobs;

/// <summary>
/// Generates WebP thumbnails for assets missing them, with wrong dimensions, or wrong format.
/// Master creates one sub-job per batch of assets.
/// </summary>
public class ThumbnailJob : Job {
    ThumbnailJob?         ParentJob;
    Asset[]?              Batch;
    int                   ThumbnailSize;
    int                   ThumbnailQuality;
    string?               ThumbnailPath;
    string                ExpectedFormat = "webp";
    int                   processedAssets;
    int                   failedAssets;
    int                   totalAssets = 1;
    int                   batchSize;
    List<Job>             childJobs = [];
    ILogger<ThumbnailJob> Logger;
    ISettingsRepository   SettingsRepository;
    IAssetRepository      AssetRepository;
    JobManager?           JobManager;

    /// <inheritdoc />
    public override string Name { get; }

    /// <summary>
    /// Creates a master thumbnail generation job.
    /// </summary>
    public ThumbnailJob(
        ILogger<ThumbnailJob> logger,
        ISettingsRepository settingsRepository,
        IAssetRepository assetRepository,
        JobManager jobManager
    ) {
        Logger             = logger;
        SettingsRepository = settingsRepository;
        AssetRepository    = assetRepository;
        JobManager         = jobManager;
        Name               = "Thumbnail Generation Job";
    }

    /// <summary>
    /// Creates a sub-job to generate thumbnails for a batch of assets.
    /// </summary>
    public ThumbnailJob(
        ThumbnailJob parentJob,
        IEnumerable<Asset> batch,
        int thumbnailSize,
        string thumbnailPath,
        int thumbnailQuality,
        string expectedFormat,
        int batchSize,
        ILogger<ThumbnailJob> logger,
        ISettingsRepository settingsRepository,
        IAssetRepository assetRepository
    ) {
        Logger             = logger;
        SettingsRepository = settingsRepository;
        AssetRepository    = assetRepository;
        Batch              = batch.ToArray();
        Name               = $"Thumbnail Batch ({Batch.Length} assets)";
        ParentJob          = parentJob;
        ParentJobId        = parentJob.Id;
        ThumbnailSize      = thumbnailSize;
        ThumbnailPath      = thumbnailPath;
        ThumbnailQuality   = thumbnailQuality;
        ExpectedFormat     = expectedFormat;
        this.batchSize     = batchSize;
        totalAssets        = Batch.Length;
    }

    /// <inheritdoc />
    protected override async Task TaskJob(CancellationToken token) {
        if (ParentJob == null) {
            await MasterJob(token);
        } else {
            SlaveJob(token);
        }
    }

    void SlaveJob(CancellationToken token) {
        if (token.IsCancellationRequested) {
            JobStatus.Cancel("Cancellation requested from user.");
            return;
        }

        if (Batch == null || Batch.Length == 0) {
            Logger.LogWarning("Sub-job started with an empty batch. Completing.");
            JobStatus.Complete("Empty batch — nothing to process.");
            return;
        }

        JobStatus.Start();

        if (ThumbnailPath == null) {
            Logger.LogError("Thumbnail path is not set. Cannot save thumbnails.");
            JobStatus.Fail("Thumbnail path is not set.");
            return;
        }

        for (int i = 0; i < Batch.Length; i++) {
            if (token.IsCancellationRequested) {
                JobStatus.Cancel("Cancellation requested from user.");
                return;
            }

            var asset = Batch[i];
            JobStatus.UpdateProgress((float)i / Batch.Length, $"Processing {asset.OriginalPath}");

            try {
                if (string.IsNullOrEmpty(asset.OriginalPath) || !File.Exists(asset.OriginalPath)) {
                    var msg = $"Original file for asset ID {asset.Id} not found at path {asset.OriginalPath}";
                    Logger.LogWarning(msg);
                    failedAssets++;
                    continue;
                }

                if (!string.IsNullOrEmpty(asset.ThumbnailPath) && File.Exists(asset.ThumbnailPath)) {
                    Logger.LogInformation("Removing stale thumbnail for asset ID {Id} at {Path}.", asset.Id, asset.ThumbnailPath);
                    File.Delete(asset.ThumbnailPath);
                    asset.ThumbnailPath = "";
                }

                using var image = Image.Load(asset.OriginalPath);

                image.Mutate(x => x.Resize(
                        new ResizeOptions {
                            Size    = new Size(ThumbnailSize),
                            Mode    = ResizeMode.Max,
                            Sampler = KnownResamplers.CatmullRom
                        }
                    )
                );

                var path = PathFromGuid(asset.Id, ThumbnailPath!);

                if (!Directory.Exists(Path.GetDirectoryName(path)))
                    Directory.CreateDirectory(Path.GetDirectoryName(path)!);

                image.SaveAsWebp(
                    path,
                    new WebpEncoder() {
                        Quality    = ThumbnailQuality,
                        SkipMetadata = true
                    }
                );

                asset.ThumbnailPath   = path;
                asset.ThumbnailSize   = ThumbnailSize;
                asset.ThumbnailFormat = ExpectedFormat;
                processedAssets++;
            } catch (Exception ex) {
                Logger.LogError(ex, $"Failed to generate thumbnail for asset ID {asset.Id}");
                failedAssets++;
            }
        }

        AssetRepository.UpdateBulk(Batch);
        AssetRepository.Save();

        var resultMsg = $"Processed {processedAssets:N0}/{Batch.Length:N0} assets, {failedAssets:N0} failed.";
        Logger.LogInformation(resultMsg);

        if (failedAssets == Batch.Length)
            JobStatus.Fail($"All {failedAssets:N0} assets in batch failed.");
        else if (failedAssets > 0)
            JobStatus.CompleteWithErrors(resultMsg);
        else
            JobStatus.Complete(resultMsg);
    }

    async Task MasterJob(CancellationToken token) {
        if (JobManager == null) {
            Logger.LogError("JobManager is not available. Cannot create sub-jobs.");
            JobStatus.Fail("JobManager is not available. Cannot create sub-jobs.");
            return;
        }

        JobStatus.Start();
        Logger.LogInformation("Starting master Thumbnail job.");

        var thumbPathSetting    = SettingsRepository.Get(Settings.ThumbnailPath.AsString());
        var thumbSizeSetting    = SettingsRepository.Get(Settings.ThumbnailSize.AsString());
        var thumbQualitySetting = SettingsRepository.Get(Settings.ThumbnailQuality.AsString());
        var batchSizeSetting    = SettingsRepository.Get(Settings.JobBatchSize.AsString());

        if (thumbPathSetting == null) {
            Logger.LogError("Thumbnail path not found. Cannot proceed.");
            JobStatus.Fail("Thumbnail path not found.");
            return;
        }

        if (!int.TryParse(thumbSizeSetting?.Value, out var thumbnailSize)) {
            thumbnailSize = 256;
            Logger.LogWarning("Thumbnail Size setting not found or invalid. Defaulting to {Size}.", thumbnailSize);
        }

        if (!int.TryParse(thumbQualitySetting?.Value, out var thumbnailQuality)) {
            thumbnailQuality = 75;
            Logger.LogWarning("Thumbnail Quality setting not found or invalid. Defaulting to {Quality}.", thumbnailQuality);
        }

        thumbnailQuality = Math.Clamp(thumbnailQuality, 1, 100);

        batchSize = 200;
        if (batchSizeSetting != null && int.TryParse(batchSizeSetting.Value, out var parsed))
            batchSize = Math.Max(1, parsed);

        totalAssets = AssetRepository.CountAssetsMissingOrWrongThumbnail(thumbnailSize, ExpectedFormat);
        Logger.LogInformation("Found {Count} assets needing thumbnails.", totalAssets);

        if (totalAssets == 0) {
            JobStatus.Complete("No assets need thumbnails.");
            return;
        }

        var childTcs = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously);
        int pending = 0;

        for (int offset = 0; offset < totalAssets && !token.IsCancellationRequested; offset += batchSize) {
            var batch = AssetRepository.GetAssetsMissingOrWrongThumbnail(thumbnailSize, ExpectedFormat, batchSize, offset);
            var batchList = batch.ToList();

            if (batchList.Count == 0) continue;

            var job = JobManager.CreateJob<ThumbnailJob>(
                this, batchList, thumbnailSize, thumbPathSetting.Value!, thumbnailQuality, ExpectedFormat, batchSize
            );
            job.LinkParentToken(token);
            lock (childJobs) { childJobs.Add(job); }

            int capturedOffset = offset;
            job.Done += (o, _) => {
                var sub = (ThumbnailJob)o!;
                lock (childJobs) { childJobs.Remove(sub); }
                Interlocked.Add(ref processedAssets, sub.processedAssets);
                Interlocked.Add(ref failedAssets, sub.failedAssets);
                if (Interlocked.Decrement(ref pending) == 0)
                    childTcs.TrySetResult();
            };

            Interlocked.Increment(ref pending);
            JobManager.EnqueueJob(job);

            JobStatus.UpdateProgress(
                (float)offset / totalAssets,
                $"Created batch at {offset:N0}/{totalAssets:N0}."
            );
        }

        if (token.IsCancellationRequested) {
            JobStatus.Cancel("Cancellation requested from user.");
            return;
        }

        var remain = Interlocked.CompareExchange(ref pending, 0, 0);
        if (remain == 0) {
            childTcs.TrySetResult();
        } else {
            using var ctr = token.Register(() => childTcs.TrySetCanceled(token));
            try {
                await childTcs.Task;
            } catch (TaskCanceledException) {
                JobStatus.Cancel("Cancellation requested from user.");
                return;
            }
        }

        if (failedAssets > 0)
            JobStatus.CompleteWithErrors($"{failedAssets:N0} of {totalAssets:N0} thumbnails failed.");
        else
            JobStatus.Complete("All thumbnails generated successfully.");
    }

    static string PathFromGuid(Guid id, string root) {
        var s = id.ToString("N");

        return Path.Combine(
            root,
            s[..2],
            s.Substring(2, 2),
            s.Substring(4, 2),
            $"{s}.webp"
        );
    }
}