MilkyShots/Lactose/Jobs/PHashJob.cs

using Butter.Settings;
using CoenM.ImageHash.HashAlgorithms;
using Lactose.Models;
using Lactose.Repositories;
using Lactose.Utils;
using SixLabors.ImageSharp;
using SixLabors.ImageSharp.PixelFormats;
using System.Collections;

namespace Lactose.Jobs;

/// <summary>
/// Computes perceptual hashes (pHash) for assets missing them.
/// Master creates one sub-job per batch of assets.
/// </summary>
public sealed class PHashJob : Job {
    PHashJob?         ParentJob;
    Asset[]?          Batch;
    IAssetRepository  AssetRepository;
    ISettingsRepository? SettingsRepository;
    JobManager?       JobManager;
    int               processedAssets;
    int               failedAssets;
    int               totalAssets = 1;
    int               batchSize;
    List<Job>         childJobs = [];
    ILogger<PHashJob> Logger { get; }
    /// <inheritdoc />
    public override string Name { get; }

    /// <summary>
    /// Creates a master pHash job.
    /// </summary>
    public PHashJob(ILogger<PHashJob> logger, IAssetRepository assetRepository, ISettingsRepository settingsRepository, JobManager jobManager) {
        Name             = "Image Hashing";
        Logger           = logger;
        AssetRepository  = assetRepository;
        SettingsRepository = settingsRepository;
        JobManager       = jobManager;
    }

    /// <summary>
    /// Creates a sub-job to compute the pHash for a batch of assets.
    /// </summary>
    public PHashJob(PHashJob parentJob, IEnumerable<Asset> batch, int batchSize, ILogger<PHashJob> logger, IAssetRepository assetRepository) {
        ParentJob       = parentJob;
        ParentJobId     = parentJob.Id;
        Batch           = batch.ToArray();
        Name            = $"pHash Batch ({Batch.Length} assets)";
        Logger          = logger;
        AssetRepository = assetRepository;
        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();

        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 {
                PerceptualHash pHash = new();
                using var image = Image.Load<Rgba32>(asset.OriginalPath);
                asset.Hash = pHash.Hash(image).ToBitArray();
                processedAssets++;
            } catch (ArgumentNullException ex) {
                Logger.LogError(ex, $"Image at path {asset.OriginalPath} could not be found. Setting hash to 0.");
                asset.Hash = new BitArray(64);
                failedAssets++;
            } catch (UnknownImageFormatException ex) {
                Logger.LogError(ex, $"Image at path {asset.OriginalPath} is in an unknown or unsupported format.");
                asset.Hash = new BitArray(64);
                failedAssets++;
            } catch (InvalidImageContentException ex) {
                Logger.LogError(ex, $"Image at path {asset.OriginalPath} is corrupted or unreadable.");
                asset.Hash = new BitArray(64);
                failedAssets++;
            } catch (Exception ex) {
                Logger.LogError(ex, $"Failed to calculate pHash for asset {asset.Id} | {asset.OriginalPath}.");
                asset.Hash = new BitArray(64);
                failedAssets++;
            }
        }

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

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

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

    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();

        batchSize = 200;
        var batchSizeSetting = SettingsRepository?.Get(Settings.JobBatchSize.AsString());
        if (batchSizeSetting != null && int.TryParse(batchSizeSetting.Value, out var parsed))
            batchSize = Math.Max(1, parsed);

        totalAssets = AssetRepository.CountAssetsMissingPHash();
        Logger.LogInformation("Found {Count} assets missing pHash.", totalAssets);

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

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

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

            if (batchList.Count == 0) continue;

            var job = JobManager.CreateJob<PHashJob>(this, batchList, batchSize);
            job.LinkParentToken(token);
            lock (childJobs) { childJobs.Add(job); }

            job.Done += (o, _) => {
                var sub = (PHashJob)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} pHash calculations failed.");
        else
            JobStatus.Complete("All pHash calculations completed.");
    }
}