MilkyShots/Lactose/Jobs/MetadataJob.cs

using Butter.Settings;
using Lactose.Models;
using Lactose.Repositories;
using SixLabors.ImageSharp;

namespace Lactose.Jobs;

/// <summary>
/// Extracts resolution metadata for image assets missing it.
/// Master creates one sub-job per batch of assets.
/// </summary>
public sealed class MetadataJob : Job {
    MetadataJob?         ParentJob;
    Asset[]?             Batch;
    IAssetRepository     AssetRepository;
    ISettingsRepository? SettingsRepository;
    JobManager?          JobManager;
    ILogger<MetadataJob> Logger;
    int                  processedAssets;
    int                  failedAssets;
    int                  totalAssets = 1;
    int                  batchSize;
    List<Job>            childJobs = [];

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

    /// <summary>
    /// Creates a master metadata extraction job.
    /// </summary>
    public MetadataJob(ILogger<MetadataJob> logger, IAssetRepository assetRepository, ISettingsRepository settingsRepository, JobManager jobManager) {
        Logger              = logger;
        AssetRepository     = assetRepository;
        SettingsRepository  = settingsRepository;
        JobManager          = jobManager;
        Name                = "Metadata Extraction Job";
    }

    /// <summary>
    /// Creates a sub-job to extract metadata for a batch of assets.
    /// </summary>
    public MetadataJob(MetadataJob parentJob, IEnumerable<Asset> batch, int batchSize, ILogger<MetadataJob> logger, IAssetRepository assetRepository) {
        ParentJob       = parentJob;
        ParentJobId     = parentJob.Id;
        Batch           = batch.ToArray();
        Name            = $"Metadata 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 {
                if (string.IsNullOrEmpty(asset.OriginalPath) || !File.Exists(asset.OriginalPath)) {
                    Logger.LogWarning($"Original file for asset ID {asset.Id} not found at path {asset.OriginalPath}");
                    failedAssets++;
                    continue;
                }

                using var image = Image.Load(asset.OriginalPath);
                asset.ResolutionWidth  = image.Width;
                asset.ResolutionHeight = image.Height;
                processedAssets++;
            } catch (Exception ex) {
                Logger.LogError(ex, $"Failed to extract metadata for asset ID {asset.Id}");
                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.");
            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.CountAssetsMissingMetadata();
        Logger.LogInformation("Found {Count} image assets missing metadata.", totalAssets);

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

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

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

            if (batchList.Count == 0) continue;

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

            job.Done += (o, _) => {
                var sub = (MetadataJob)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} metadata extractions failed.");
        else
            JobStatus.Complete("All metadata extracted successfully.");
    }
}