MilkyShots/Lactose/Repositories/AssetRepository.cs

using Butter.Types;
using Lactose.Context;
using Lactose.Models;
using Pgvector.EntityFrameworkCore;
using System.Collections;
using System.Data;

namespace Lactose.Repositories;

/// <inheritdoc />
public class AssetRepository(LactoseDbContext context) : 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
            .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 && 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.Where(x => x.CreatedAt >= from && x.UpdatedAt <= to);

    /// <inheritdoc />
    public IEnumerable<Asset> FindByDateRange(DateTime from, DateTime to, int pageNumber, int pageSize) => 
        context.Assets.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.Where(a => a.UploadedBy == uploaderId);

    /// <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.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);
    }

    /// <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)
            .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);
        return context.Assets.Where(a => (a.ThumbnailPath == null || a.ThumbnailPath == "") && a.DeletedAt == null);
    }

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

    /// <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)
            .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);

    /// <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)
            .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);

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

    /// <inheritdoc />
    public IEnumerable<Asset> 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) {
        var query = context.Assets.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);

        // 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 && a.UploadedBy == userId)
            ))
        };

        total = query.Count();

        if (orderRandomly && seed.HasValue) {
            var allIds = query.Select(a => a.Id).ToList();
            var seedHash = seed.Value.GetHashCode();
            var shuffledIds = allIds
                .OrderBy(id => DeterministicHash(id, seedHash))
                .Skip(pageNumber * pageSize)
                .Take(pageSize)
                .ToList();

            var assets = context.Assets
                .AsSplitQuery()
                .Where(a => shuffledIds.Contains(a.Id))
                .Include(a => a.Albums!).ThenInclude(a => a.PersonOwner)
                .ToList();

            return shuffledIds.Select(id => assets.First(a => a.Id == id)).ToList();
        }

        IQueryable<Asset> dataQuery = query
            .AsSplitQuery()
            .Include(a => a.Albums!).ThenInclude(a => a.PersonOwner);
        dataQuery = orderRandomly ? dataQuery.OrderBy(a => a.Id) : dataQuery.OrderByDescending(a => a.CreatedAt);

        return dataQuery
            .Skip(pageNumber * pageSize)
            .Take(pageSize)
            .ToList();
    }

    private static int DeterministicHash(Guid id, int seed) {
        var bytes = id.ToByteArray();
        unchecked {
            var hash = seed;
            for (var i = 0; i < bytes.Length; i++)
                hash = hash * 31 + bytes[i];
            return hash & 0x7FFFFFFF;
        }
    }

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

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