MilkyShots/Lactose/Repositories/AssetRepository.cs

using Butter.Dtos.Asset;
using Butter.Types;
using Lactose.Context;
using Lactose.Mapper;
using Lactose.Models;
using Microsoft.EntityFrameworkCore;
using Microsoft.Extensions.Logging;
using Pgvector.EntityFrameworkCore;
using System.Collections;
using System.Data;

namespace Lactose.Repositories;

/// <inheritdoc />
public class AssetRepository(LactoseDbContext context, ILogger<AssetRepository> logger) : IAssetRepository {
    
    /// <inheritdoc />
    public Asset? Find(Guid id) => context.Assets.Find(id);

    /// <inheritdoc />
    public Asset? FindVisible(Guid id, Guid? userId, EAccessLevel accessLevel) {
        IQueryable<Asset> query = context.Assets
            .AsNoTracking()
            .AsSplitQuery()
            .Include(a => a.Uploader)
            .Where(a => a.Id == id);

        return accessLevel switch {
            >= EAccessLevel.Admin => query.FirstOrDefault(),
            EAccessLevel.Curator => query.FirstOrDefault(a =>
                a.DeletedAt == null || a.UploadedBy == userId),
            EAccessLevel.Maintainer when userId.HasValue => query
                .Include(a => a.Albums)
                .FirstOrDefault(a => a.DeletedAt == null && (
                    a.Visibility == EVisibility.Public ||
                    a.Visibility == EVisibility.Protected ||
                    a.UploadedBy == userId ||
                    a.Albums!.Any(al => al.PersonOwnerId.HasValue &&
                        context.PersonMaintainers.Any(pm =>
                            pm.UserId == userId.Value && pm.PersonId == al.PersonOwnerId.Value))
                )),
            _ => query.FirstOrDefault(a =>
                a.DeletedAt == null && (
                    a.Visibility == EVisibility.Public ||
                    (a.Visibility == EVisibility.Protected && userId.HasValue) ||
                (a.Visibility == EVisibility.Private && userId.HasValue && a.UploadedBy == userId)
                ))
        };
    }

    /// <inheritdoc />
    public Asset? FindWithAlbums(Guid id) => context.Assets
        .AsSplitQuery()
        .Include(a => a.Albums)
        .FirstOrDefault(a => a.Id == id);

    /// <inheritdoc />
    public IEnumerable<Asset> FindByDateRange(DateTime from, DateTime to) => 
        context.Assets.AsNoTracking().Where(x => x.CreatedAt >= from && x.UpdatedAt <= to);

    /// <inheritdoc />
    public IEnumerable<Asset> FindByDateRange(DateTime from, DateTime to, int pageNumber, int pageSize) => 
        context.Assets.AsNoTracking().Where(x => x.CreatedAt >= from && x.UpdatedAt <= to)
            .Skip(pageNumber * pageSize)
            .Take(pageSize);
    
    /// <inheritdoc />
    public void Save() => context.SaveChanges();

    /// <inheritdoc />
    public void Insert(Asset asset) {
        if(!context.Assets.Any(x => x.OriginalPath == asset.OriginalPath))
            context.Assets.Add(asset);
        else
            throw new DuplicateNameException("The required asset already exists in the database.");
    }

    /// <inheritdoc />
    public void Update(Asset asset) {
        asset.UpdatedAt = DateTime.UtcNow;
        context.Assets.Update(asset);
    }

    /// <inheritdoc />
    public void UpdateBulk(IEnumerable<Asset> assets) {
        foreach (var asset in assets) {
            asset.UpdatedAt = DateTime.UtcNow;
            context.Assets.Update(asset);
        }
    }

    /// <inheritdoc />
    public IEnumerable<Asset> FindByUploader(Guid uploaderId) => context.Assets.AsNoTracking().Where(a => a.UploadedBy == uploaderId);

    /// <inheritdoc />
    public int BulkSetVisibilityByAlbumIds(IEnumerable<Guid> albumIds, EVisibility visibility, Guid? userId, EAccessLevel accessLevel) {
        // Scope the update to assets the caller is permitted to modify (R2). Curators/admins
        // see all non-deleted assets; maintainers additionally respect their uploads and
        // the albums of people they maintain.
        var query = context.Assets
            .Where(a => a.DeletedAt == null && a.Albums!.Any(al => albumIds.Contains(al.Id)));

        query = accessLevel switch {
            EAccessLevel.Admin or EAccessLevel.Curator => query,
            EAccessLevel.Maintainer when userId.HasValue => query.Where(a =>
                a.Visibility == EVisibility.Public ||
                a.Visibility == EVisibility.Protected ||
                a.UploadedBy == userId ||
                a.Albums!.Any(al => al.PersonOwnerId.HasValue &&
                    context.PersonMaintainers.Any(pm =>
                        pm.UserId == userId && pm.PersonId == al.PersonOwnerId.Value))),
            _ => query.Where(a => false)
        };

        return query.ExecuteUpdate(s => s
            .SetProperty(a => a.Visibility, visibility)
            .SetProperty(a => a.UpdatedAt, DateTime.UtcNow));
    }

    /// <inheritdoc />
    public int BulkSetVisibilityByPersonIds(IEnumerable<Guid> personIds, EVisibility visibility) =>
        context.Assets
            .Where(a => a.DeletedAt == null && a.Albums!.Any(al =>
                al.PersonOwnerId.HasValue && personIds.Contains(al.PersonOwnerId.Value)))
            .ExecuteUpdate(s => s
                .SetProperty(a => a.Visibility, visibility)
                .SetProperty(a => a.UpdatedAt, DateTime.UtcNow));

    /// <inheritdoc />
    public IEnumerable<Asset> FindBulk(IEnumerable<Guid> ids) => context.Assets
        .AsSplitQuery()
        .Include(a => a.Albums)
        .Where(a => ids.Contains(a.Id));
    
    /// <inheritdoc />
    public Asset? FindByPath(string path) => context.Assets.AsNoTracking().FirstOrDefault(a => a.OriginalPath == path);

    /// <inheritdoc />
    public int CountAssetsMissingPHash() {
        var empty = new BitArray(64);
        return context.Assets.Count(a => a.Hash == empty && a.DeletedAt == null && a.ProcessFailedAt == null);
    }

    /// <inheritdoc />
    public IEnumerable<Asset> GetAssetsMissingPHash(int limit, int offset) {
        var empty = new BitArray(64);
        return context.Assets.Where(a => a.Hash == empty && a.DeletedAt == null && a.ProcessFailedAt == null)
            .OrderBy(a => a.Id)
            .Skip(offset)
            .Take(limit);
    }
    
    /// <inheritdoc />
    public IEnumerable<List<Asset>> GetDuplicates() => context.Assets
        .GroupBy(a => new { a.Hash })
        .Where(g => g.Count() > 1)
        .Select(g => g.ToList());
    
    /// <inheritdoc />
    public IEnumerable<Asset> GetWithinHammingDistance(ulong phash, int distance) {
        if(distance < 0 || distance > 63)
            throw new ArgumentOutOfRangeException(nameof(distance), "Distance must be between 0 and 64.");
        
        return context.Assets
            .OrderBy(a => a.HammingDistance(phash))
            .Where(a => a.Hash.HammingDistance(phash) <= distance);
    }

    /// <inheritdoc />
    public IEnumerable<Asset> GetWithinHammingDistance(ulong phash, float distance) 
        => GetWithinHammingDistance(phash,  (int)(distance * 64));

    /// <inheritdoc />
    public IEnumerable<Asset> GetAssetsMissingThumbnail(out int count) {
        count = context.Assets.Count(a => (a.ThumbnailPath == null || a.ThumbnailPath == "") && a.DeletedAt == null && a.ProcessFailedAt == null && a.MimeType != "image/gif");
        return context.Assets.Where(a => (a.ThumbnailPath == null || a.ThumbnailPath == "") && a.DeletedAt == null && a.ProcessFailedAt == null && a.MimeType != "image/gif");
    }

    /// <inheritdoc />
    public int CountAssetsMissingOrWrongThumbnail(int thumbnailSize, string expectedFormat) =>
        context.Assets.Count(a => (a.ThumbnailPath == null || a.ThumbnailPath == "" || a.ThumbnailSize != thumbnailSize || a.ThumbnailFormat != expectedFormat) && a.DeletedAt == null && a.ProcessFailedAt == null && a.MimeType != "image/gif");

    /// <inheritdoc />
    public IEnumerable<Asset> GetAssetsMissingOrWrongThumbnail(int thumbnailSize, string expectedFormat, int limit, int offset) =>
        context.Assets
            .Where(a => (a.ThumbnailPath == null || a.ThumbnailPath == "" || a.ThumbnailSize != thumbnailSize || a.ThumbnailFormat != expectedFormat) && a.DeletedAt == null && a.ProcessFailedAt == null && a.MimeType != "image/gif")
            .OrderBy(a => a.Id)
            .Skip(offset)
            .Take(limit);

    /// <inheritdoc />
    public int CountAssetsMissingOrWrongPreview(int previewSize, string expectedFormat) =>
        context.Assets.Count(a => (a.PreviewPath == null || a.PreviewPath == "" || a.PreviewSize != previewSize || a.PreviewFormat != expectedFormat) && a.DeletedAt == null && a.ProcessFailedAt == null && a.MimeType != "image/gif");

    /// <inheritdoc />
    public IEnumerable<Asset> GetAssetsMissingOrWrongPreview(int previewSize, string expectedFormat, int limit, int offset) =>
        context.Assets
            .Where(a => (a.PreviewPath == null || a.PreviewPath == "" || a.PreviewSize != previewSize || a.PreviewFormat != expectedFormat) && a.DeletedAt == null && a.ProcessFailedAt == null && a.MimeType != "image/gif")
            .OrderBy(a => a.Id)
            .Skip(offset)
            .Take(limit);
    
    /// <inheritdoc />
    public int CountAssetsMissingMetadata() =>
        context.Assets.Count(a => a.Type == EAssetType.Image && a.ResolutionWidth == 0 && a.DeletedAt == null && a.ProcessFailedAt == null);

    /// <inheritdoc />
    public IEnumerable<Asset> GetAssetsMissingMetadata(int limit, int offset) =>
        context.Assets.Where(a => a.Type == EAssetType.Image && a.ResolutionWidth == 0 && a.DeletedAt == null && a.ProcessFailedAt == null)
            .OrderBy(a => a.Id)
            .Skip(offset)
            .Take(limit);

    /// <inheritdoc />
    public int CountAssetsNeedingConversion() =>
        context.Assets.Count(a => a.MimeType == "image/gif" && a.DeletedAt == null && a.ProcessFailedAt == null && (
            a.ConvertedPath == null || a.ConvertedPath == "" ||
            a.ThumbnailPath == null || a.ThumbnailPath == ""));

    /// <inheritdoc />
    public IEnumerable<Asset> GetAssetsNeedingConversion(int limit, int offset) =>
        context.Assets
            .Where(a => a.MimeType == "image/gif" && a.DeletedAt == null && a.ProcessFailedAt == null && (
                a.ConvertedPath == null || a.ConvertedPath == "" ||
                a.ThumbnailPath == null || a.ThumbnailPath == ""))
            .OrderBy(a => a.Id)
            .Skip(offset)
            .Take(limit);

    /// <inheritdoc />
    public int CountBrokenAssets() =>
        context.Assets.Count(a => a.ProcessFailedAt != null && a.DeletedAt == null);

    /// <inheritdoc />
    public IEnumerable<AssetPreviewDto> GetBrokenAssets(int page, int pageSize, out int total) {
        total = CountBrokenAssets();

        var assets = context.Assets
            .AsNoTracking()
            .AsSplitQuery()
            .Include(a => a.Albums!).ThenInclude(al => al.PersonOwner)
            .Where(a => a.ProcessFailedAt != null && a.DeletedAt == null)
            .OrderBy(a => a.ProcessFailedAt)
            .Skip(page * pageSize)
            .Take(pageSize)
            .ToList();

        return assets.Select(a => a.ToAssetPreviewDto(EAccessLevel.Admin, null)).ToList();
    }

    /// <inheritdoc />
    public void ClearProcessError(IEnumerable<Guid> ids) {
        var assets = context.Assets.Where(a => ids.Contains(a.Id)).ToList();
        foreach (var asset in assets) {
            asset.ProcessFailedAt = null;
            asset.ProcessErrorMessage = null;
        }
    }

    /// <inheritdoc />
    public IEnumerable<AssetPreviewDto> GetAssets(EAssetType? type, DateTime? from, DateTime? to, bool orderRandomly, Guid? seed, int pageNumber, int pageSize, out int total, Guid? userId = null, EAccessLevel accessLevel = EAccessLevel.User, bool unlinked = false, Guid? folderId = null, Guid? uploadedBy = null, string? search = null, bool includeCount = true) {
        var query = context.Assets.AsNoTracking().AsQueryable();
        if (type.HasValue)
            query = query.Where(a => a.Type == type.Value);
        if (from.HasValue)
            query = query.Where(a => a.CreatedAt >= from.Value);
        if (to.HasValue)
            query = query.Where(a => a.UpdatedAt <= to.Value);
        if (unlinked)
            query = query.Where(a => a.Albums!.Count == 0);
        if (folderId.HasValue)
            query = query.Where(a => a.FolderId == folderId.Value);
        if (uploadedBy.HasValue)
            query = query.Where(a => a.UploadedBy == uploadedBy.Value);
        if (!string.IsNullOrEmpty(search))
            query = query.Where(a => EF.Functions.ILike(a.OriginalFilename, $"%{search}%"));

        // Apply visibility filter
        query = accessLevel switch {
            EAccessLevel.Admin => query,
            EAccessLevel.Curator => query.Where(a => a.DeletedAt == null || a.UploadedBy == userId),
            EAccessLevel.Maintainer => query.Where(a => a.DeletedAt == null && (
                a.Visibility == EVisibility.Public ||
                a.Visibility == EVisibility.Protected ||
                a.UploadedBy == userId ||
                a.Albums!.Any(al => al.PersonOwnerId.HasValue &&
                    context.PersonMaintainers.Any(pm =>
                        pm.UserId == userId && pm.PersonId == al.PersonOwnerId.Value))
            )),
            _ => query.Where(a => a.DeletedAt == null && (
                a.Visibility == EVisibility.Public ||
                (a.Visibility == EVisibility.Protected && userId.HasValue) ||
                (a.Visibility == EVisibility.Private && userId.HasValue && a.UploadedBy == userId)
            ))
        };

        total = includeCount ? query.Count() : 0;

        // Deterministic page of IDs first: random order is a seeded md5 shuffle on the server,
        // so the full ID set is never materialized in memory.
        IQueryable<Asset> ordered = orderRandomly switch {
            true when seed.HasValue => query.OrderBy(a => PgFunctions.Md5(a.Id.ToString() + seed.Value.ToString())),
            true => query.OrderBy(a => a.Id),
            false => query.OrderByDescending(a => a.CreatedAt)
        };

        var pageIds = ordered
            .Skip(pageNumber * pageSize)
            .Take(pageSize)
            .Select(a => a.Id)
            .ToList();

        if (pageIds.Count == 0)
            return [];

        var assets = query
            .Where(a => pageIds.Contains(a.Id))
            .Select(a => new AssetPreviewDto {
                Id               = a.Id,
                MimeType         = a.ConvertedMimeType != null && a.ConvertedMimeType != "" ? a.ConvertedMimeType : a.MimeType,
                ResolutionWidth  = a.ResolutionWidth,
                ResolutionHeight = a.ResolutionHeight,
                HasThumbnail     = a.ThumbnailPath != null && a.ThumbnailPath != "",
                HasPreview       = a.PreviewPath != null && a.PreviewPath != "",
                FileName         = accessLevel >= EAccessLevel.Curator ? a.OriginalFilename : null,
                Visibility       = accessLevel >= EAccessLevel.Maintainer ? a.Visibility : null,
                DeletedAt        = accessLevel >= EAccessLevel.Admin || a.UploadedBy == userId ? a.DeletedAt : null
            })
            .ToList();

        // Album and cosplayer names per asset via a single grouped query
        var albumLinks = (from a in context.Assets
                          from al in a.Albums!
                          where pageIds.Contains(a.Id)
                          select new { AssetId = a.Id, al.Id, al.Title, al.PersonOwnerId, PersonName = al.PersonOwner!.Name })
            .ToList();

        var linksByAsset = albumLinks
            .GroupBy(x => x.AssetId)
            .ToDictionary(g => g.Key, g => g.ToList());

        foreach (var dto in assets) {
            if (!linksByAsset.TryGetValue(dto.Id, out var links)) continue;
            dto.AlbumNames       = links.Select(x => x.Title).Distinct().ToList();
            dto.AlbumIds         = links.Select(x => x.Id).Distinct().ToList();
            dto.CosplayerNames   = links.Where(x => x.PersonOwnerId.HasValue).Select(x => x.PersonName).Distinct().ToList();
            dto.CosplayerIds     = links.Where(x => x.PersonOwnerId.HasValue).Select(x => x.PersonOwnerId!.Value).Distinct().ToList();
        }

        var orderMap = pageIds.Select((id, i) => (id, i)).ToDictionary(x => x.id, x => x.i);
        return [.. assets.OrderBy(a => orderMap.GetValueOrDefault(a.Id))];
    }

    /// <inheritdoc />
    public void Dispose() => context.Dispose();

    /// <inheritdoc />
    public List<string> GetRandomPathsByFolder(Guid folderId, int count) =>
        context.Assets.AsNoTracking()
            .Where(a => a.FolderId == folderId && a.DeletedAt == null)
            .OrderBy(_ => Guid.NewGuid())
            .Take(count)
            .Select(a => a.OriginalPath)
            .ToList();

    /// <inheritdoc />
    public List<Butter.Dtos.Asset.AssetGroupDto> GetFolderGroups(Guid? userId, EAccessLevel accessLevel, bool unlinkedOnly = true) {
        logger.LogTrace("GetFolderGroups: userId={UserId}, accessLevel={AccessLevel}, unlinkedOnly={UnlinkedOnly}", userId, accessLevel, unlinkedOnly);

        var query = context.Assets.AsNoTracking()
            .Where(a => a.DeletedAt == null);

        if (unlinkedOnly)
            query = query.Where(a => a.Albums!.Count == 0);

        if (accessLevel < EAccessLevel.Curator) {
            if (accessLevel == EAccessLevel.Maintainer && userId.HasValue) {
                var count = query.Count(a => a.UploadedBy == userId);
                return [new Butter.Dtos.Asset.AssetGroupDto {
                    FolderId = null,
                    Name = "My Uploads",
                    AssetCount = count
                }];
            }
            return [];
        }

        var groups = query
            .GroupBy(a => a.FolderId)
            .Select(g => new {
                FolderId = g.Key,
                AssetCount = g.Count()
            })
            .ToList();

        var folderNames = context.Folders
            .Where(f => groups.Select(g => g.FolderId).Contains(f.Id))
            .ToDictionary(f => f.Id, f => f.BasePath);

        var result = groups
            .Where(g => g.AssetCount > 0)
            .Select(g => new Butter.Dtos.Asset.AssetGroupDto {
                FolderId = g.FolderId,
                Name = g.FolderId.HasValue && folderNames.TryGetValue(g.FolderId.Value, out var name) ? name : "No Folder",
                AssetCount = g.AssetCount
            })
            .OrderByDescending(g => g.AssetCount)
            .ToList();

        return result;
    }

}