MilkyShots/Lactose/Jobs/FileSystemCrawlJob.cs

using Butter;
using Butter.Settings;
using Butter.Types;
using Lactose.Context;
using Lactose.Models;
using Lactose.Repositories;
using Microsoft.EntityFrameworkCore;
using Npgsql;
using System.Collections;
using static System.String;

namespace Lactose.Jobs;

/// <summary>
/// Crawls a file system directory, incrementally discovering and reconciling assets against the database.
/// Uses per-directory modification timestamps to skip file enumeration in unchanged directories,
/// and batches all database writes instead of saving once per file.
/// </summary>
public sealed class FileSystemCrawlJob : Job {
    /// <inheritdoc />
    public override string      Name { get; }
    /// <inheritdoc />
    public override JobStatus   JobStatus { get; }
    readonly Guid                          folderId;
    readonly string                        workingPath;
    readonly bool                          deep;
    readonly ILogger<FileSystemCrawlJob>   logger;

    const int ProgressUpdateEveryDirs = 200;

    /// <summary>
    /// Initialises a new file-system crawl job for the specified folder.
    /// </summary>
    /// <param name="folderId">The folder ID.</param>
    /// <param name="workingPath">The base path to crawl.</param>
    /// <param name="deep">When true, ignores recorded scan state and reconciles every asset.</param>
    /// <param name="logger">Logger instance.</param>
    public FileSystemCrawlJob(Guid folderId, string workingPath, bool deep, ILogger<FileSystemCrawlJob> logger) {
        JobStatus            = new(this);
        this.folderId        = folderId;
        this.workingPath     = workingPath;
        this.deep            = deep;
        this.logger          = logger;
        Name                 = $"Directory Scan: {workingPath}";
    }

    ///<inheritdoc />
    protected override async Task TaskJob(CancellationToken token) {
        logger.LogInformation("Started crawling directory {Path} (deep={Deep})", workingPath, deep);

        var root = new DirectoryInfo(workingPath);
        if (!root.Exists) {
            logger.LogWarning("Directory {Dir} does not exist. Skipping crawl.", workingPath);
            JobStatus.Complete("Directory does not exist.");
            return;
        }

        var db           = Scope!.ServiceProvider.GetRequiredService<LactoseDbContext>();
        var settingsRepo = Scope.ServiceProvider.GetRequiredService<ISettingsRepository>();

        Guid? uploaderId = null;
        var uploaderSetting = settingsRepo.Get(Settings.SystemUploaderId.AsString())?.Value;
        if (Guid.TryParse(uploaderSetting, out var parsed)) uploaderId = parsed;

        // Separate context for asset reconciliation so the change tracker can be cleared per directory
        // without detaching the scan-state entities tracked by `db`.
        using var assetScope = Scope.ServiceProvider.CreateScope();
        var assetDb = assetScope.ServiceProvider.GetRequiredService<LactoseDbContext>();

        var existingStates = await db.DirectoryScanStates
            .Where(s => s.FolderId == folderId)
            .ToDictionaryAsync(s => s.Path, StringComparer.Ordinal, token);

        // Progress denominator: reuse the number of directories recorded by the previous scan, which is
        // accurate once a first run has completed. On the very first run (no state yet) there is no
        // denominator, so the bar ramps heuristically while the live message shows real progress.
        int totalDirs = Math.Max(existingStates.Count, 1);
        bool haveDenominator = existingStates.Count > 0;
        logger.LogInformation("Crawling {Path}: {Total} known directories.", workingPath, totalDirs);
        JobStatus.UpdateProgress(0.02f, haveDenominator
            ? $"Scanning {totalDirs:N0} directories..."
            : "Scanning... (first run)");

        int dirsProcessed = 0, dirsSkipped = 0, walkedDirs = 0,
            filesAdded = 0, filesUpdated = 0, filesDeleted = 0;
        var lastProgressUpdate = DateTime.UtcNow;
        var seenPaths = new HashSet<string>(StringComparer.Ordinal);

        var stack = new Stack<string>();
        stack.Push(root.FullName);

        while (stack.Count > 0) {
            token.ThrowIfCancellationRequested();
            var currentPath = stack.Pop();
            seenPaths.Add(currentPath);

            var currentDir = new DirectoryInfo(currentPath);
            if (!currentDir.Exists) continue;
            walkedDirs++;

            // A directory's mtime changes only when its direct children change; descendant changes bump
            // only their own parent. So we must always descend, but skip file reconciliation for clean dirs.
            bool changed = deep;
            if (!changed) {
                if (!existingStates.TryGetValue(currentPath, out var state)) {
                    changed = true; // new directory
                } else {
                    changed = state.LastMtime != currentDir.LastWriteTimeUtc;
                }
            }

            List<DirectoryInfo> subdirs;
            FileInfo[]? files = null;
            try {
                subdirs = currentDir.GetDirectories().ToList();
                if (changed) files = currentDir.GetFiles();
            } catch (Exception ex) {
                logger.LogWarning(ex, "Could not enumerate directory {Path}.", currentPath);
                continue;
            }

            if (changed) {
                dirsProcessed++;
                bool reconciled = await ReconcileDirectory(assetDb, currentPath, files!, uploaderId, token,
                    (added, updated, deleted) => { filesAdded += added; filesUpdated += updated; filesDeleted += deleted; });

                // Only record scan state when the reconciliation batch actually committed; a dropped
                // batch leaves the mtime stale so the directory is retried on the next run.
                if (reconciled) {
                    if (existingStates.TryGetValue(currentPath, out var existingState)) {
                        existingState.LastMtime = currentDir.LastWriteTimeUtc;
                    } else {
                        existingStates[currentPath] = new DirectoryScanState {
                            Id          = Guid.NewGuid(),
                            FolderId    = folderId,
                            Path        = currentPath,
                            LastMtime   = currentDir.LastWriteTimeUtc
                        };
                        db.DirectoryScanStates.Add(existingStates[currentPath]);
                    }
                }
            } else {
                dirsSkipped++;
            }

            foreach (var subdir in subdirs)
                stack.Push(subdir.FullName);

            if (walkedDirs % ProgressUpdateEveryDirs == 0
                || DateTime.UtcNow - lastProgressUpdate > TimeSpan.FromMilliseconds(500)) {
                lastProgressUpdate = DateTime.UtcNow;
                // Ramp heuristically when we have no accurate denominator (first run); otherwise
                // progress is walkedDirs/totalDirs. Cap before the final 0.95 "saving" step.
                float progress = haveDenominator
                    ? Math.Clamp(0.05f + 0.90f * walkedDirs / totalDirs, 0.05f, 0.90f)
                    : Math.Clamp(0.05f + 0.80f * Math.Min(1f, walkedDirs / 5000f), 0.05f, 0.85f);
                JobStatus.UpdateProgress(progress,
                    haveDenominator
                        ? $"Scanning {workingPath}: {walkedDirs:N0}/{totalDirs:N0} dirs — {filesAdded:N0} added, {filesUpdated:N0} updated, {filesDeleted:N0} removed"
                        : $"Scanning {workingPath}: {walkedDirs:N0} dirs — {filesAdded:N0} added, {filesUpdated:N0} updated, {filesDeleted:N0} removed");
            }
        }

        JobStatus.UpdateProgress(0.95f, "Saving scan state...");

        // Drop scan-state rows for directories that no longer exist on disk.
        var staleStates = existingStates.Values.Where(s => !seenPaths.Contains(s.Path)).ToList();
        db.DirectoryScanStates.RemoveRange(staleStates);

        var folderEntity = await db.Folders.FirstOrDefaultAsync(f => f.Id == folderId, token);
        if (folderEntity != null) {
            folderEntity.LastFileScanAt = DateTime.UtcNow;
        }

        await db.SaveChangesAsync(token);

        logger.LogInformation("Crawl complete for {Path}: {Dirs} dirs processed, {Skipped} skipped; {Added} added, {Updated} updated, {Deleted} removed.",
            workingPath, dirsProcessed, dirsSkipped, filesAdded, filesUpdated, filesDeleted);
        JobStatus.Complete($"Scanned {workingPath}: {filesAdded:N0} added, {filesUpdated:N0} updated, {filesDeleted:N0} removed, {dirsSkipped:N0} dirs unchanged.");
    }

    /// <summary>
    /// Reconciles the direct children of a directory against the database: inserts new assets,
    /// updates changed ones, and soft-deletes previously tracked files that no longer exist.
    /// All writes are batched into a single <see cref="DbContext.SaveChangesAsync(CancellationToken)"/> call.
    /// </summary>
    /// <param name="assetDb">A dedicated context for asset writes; cleared after each directory.</param>
    /// <param name="dirPath">The directory being reconciled.</param>
    /// <param name="files">The current direct files of the directory.</param>
    /// <param name="uploaderId">The system uploader id to backfill, if any.</param>
    /// <param name="token">Cancellation token.</param>
    /// <param name="report">Callback receiving (added, updated, deleted) counts.</param>
    /// <returns>True when the batch committed; false when it was dropped due to a duplicate-key race.</returns>
    async Task<bool> ReconcileDirectory(LactoseDbContext assetDb, string dirPath, FileInfo[] files, Guid? uploaderId,
        CancellationToken token, Action<int, int, int> report) {
        var prefix = dirPath.TrimEnd(Path.DirectorySeparatorChar) + Path.DirectorySeparatorChar;
        var lo = prefix;
        var hi = prefix + "\uFFFF";

        // Load known assets that are direct children of this directory. Soft-deleted assets are included
        // so a file that still exists on disk is reconciled (or left deleted) instead of re-inserted,
        // which would violate the unique index on OriginalPath. The path prefix (not FolderId) is what
        // scopes the lookup — OriginalPath is globally unique, so a row can exist under any folder.
        var known = (await assetDb.Assets.AsNoTracking()
                .Where(a => a.OriginalPath.CompareTo(lo) >= 0 && a.OriginalPath.CompareTo(hi) < 0)
                .ToListAsync(token))
            .Where(a => !a.OriginalPath.AsSpan(prefix.Length).Contains(Path.DirectorySeparatorChar))
            .ToDictionary(a => a.OriginalPath, StringComparer.Ordinal);

        // Exact-path fallback: OriginalPath is globally unique, so a row can exist under a different
        // folder (e.g. a recreated folder with a new ID) or fall outside the range query above.
        // Such files must be reconciled, never re-inserted.
        var unmatched = files
            .Select(f => f.FullName)
            .Where(p => !known.ContainsKey(p))
            .Distinct(StringComparer.Ordinal)
            .ToList();
        foreach (var chunk in unmatched.Chunk(1000)) {
            var existing = await assetDb.Assets.AsNoTracking()
                .Where(a => chunk.Contains(a.OriginalPath))
                .ToListAsync(token);
            foreach (var a in existing)
                known.TryAdd(a.OriginalPath, a);
        }

        int added = 0, updated = 0, deleted = 0;

        foreach (var file in files) {
            token.ThrowIfCancellationRequested();
            if (known.TryGetValue(file.FullName, out var asset)) {
                known.Remove(file.FullName);
                if (file.LastWriteTimeUtc > (asset.UpdatedAt ?? asset.CreatedAt)) {
                    assetDb.Attach(asset);
                    asset.FileSize  = file.Length;
                    asset.UpdatedAt = file.LastWriteTimeUtc;
                    if (asset.UploadedBy == null && uploaderId.HasValue) asset.UploadedBy = uploaderId;
                    updated++;
                } else if (asset.UploadedBy == null && uploaderId.HasValue) {
                    assetDb.Attach(asset);
                    asset.UploadedBy = uploaderId;
                    updated++;
                }
            } else {
                var newAsset = AssetFromPath(file.FullName, uploaderId);
                if (newAsset != null) {
                    assetDb.Assets.Add(newAsset);
                    added++;
                }
            }
        }

        foreach (var gone in known.Values) {
            token.ThrowIfCancellationRequested();
            if (gone.DeletedAt != null) continue; // already soft-deleted
            assetDb.Attach(gone);
            gone.DeletedAt = DateTime.UtcNow;
            deleted++;
        }

        if (added + updated + deleted > 0) {
            try {
                await assetDb.SaveChangesAsync(token);
            } catch (DbUpdateException ex) when (IsDuplicateKey(ex)) {
                // Last-resort safety net for genuine races (e.g. two overlapping folders inserting the
                // same new path concurrently). The batch is dropped; the file is found next run.
                logger.LogDebug(ex, "Dropped {Count} asset change(s) in {Path} due to a duplicate asset path.",
                    added + updated + deleted, dirPath);
                report(0, 0, 0);
                return false;
            } finally {
                assetDb.ChangeTracker.Clear();
            }
        }

        report(added, updated, deleted);
        return true;
    }

    static bool IsDuplicateKey(DbUpdateException ex) =>
        ex.InnerException is PostgresException { SqlState: "23505" }
        || ex.InnerException?.InnerException is PostgresException { SqlState: "23505" };

    /// <summary>
    /// Builds an <see cref="Asset"/> from a file path, or null when the extension is not a supported media type.
    /// </summary>
    /// <param name="filePath">The full path of the file.</param>
    /// <param name="uploaderId">The uploader id to assign, if any.</param>
    /// <returns>A new asset, or null if the file type is not supported.</returns>
    Asset? AssetFromPath(string filePath, Guid? uploaderId) {
        logger.LogTrace("Loading asset from file path {filePath}", filePath);
        var finfo = new FileInfo(filePath);

        string ext = finfo.Extension.ToLower();
        EAssetType? type;

        var mimeImg = MimeTypes.Image.FirstOrDefault(mime => mime.Extensions.Contains(ext), new(Empty, []));
        var mimeVid = MimeTypes.Video.FirstOrDefault(mime => mime.Extensions.Contains(ext), new(Empty, []));

        if (mimeImg.MimeType != Empty) type      = EAssetType.Image;
        else if (mimeVid.MimeType != Empty) type = EAssetType.Video;
        else return null;

        return new Asset {
            FolderId         = folderId,
            OriginalPath     = finfo.FullName,
            OriginalFilename = finfo.Name,
            Type             = type.Value,
            Visibility       = EVisibility.Private,
            UploadedBy       = uploaderId,
            CreatedAt        = finfo.CreationTimeUtc,
            UpdatedAt        = finfo.LastWriteTimeUtc,
            IngestedAt       = DateTime.UtcNow,
            MimeType         = type.Value switch {
                EAssetType.Image => mimeImg.MimeType,
                EAssetType.Video => mimeVid.MimeType,
                _                => throw new ArgumentOutOfRangeException()
            },
            FileSize = finfo.Length,
            Hash     = new BitArray(64)
        };
    }
}