MilkyShots/Lactose/Repositories/PersonRepository.cs

using Butter.Dtos;
using Butter.Dtos.Person;
using Butter.Types;
using Lactose.Context;
using Lactose.Mapper;
using Lactose.Models;
using Microsoft.EntityFrameworkCore;

namespace Lactose.Repositories;

/// <inheritdoc />
public class PersonRepository(LactoseDbContext context) : IPersonRepository {
    /// <inheritdoc />
    public void Insert(Person person) => context.People.Add(person);

    /// <inheritdoc />
    public Task SaveAsync(CancellationToken cancellationToken) => context.SaveChangesAsync(cancellationToken);

    /// <inheritdoc />
    public async Task<TResult?> ExecuteInTransactionAsync<TResult>(Func<Task<TResult?>> operation) {
        await using var tx = await context.Database.BeginTransactionAsync();
        var result = await operation();
        await tx.CommitAsync();
        return result;
    }

    /// <inheritdoc />
    public Task<Person?> FindAsync(Guid id, CancellationToken cancellationToken) =>
        context.People.FirstOrDefaultAsync(p => p.Id == id, cancellationToken);

    /// <inheritdoc />
    public async Task<Person?> FindVisibleAsync(Guid id, Guid? userId, EAccessLevel accessLevel,
        string? albumSearch, string? albumSortBy, bool albumSortAsc,
        int albumPage, int albumPageSize, CancellationToken cancellationToken) {
        var person = await context.People
            .AsNoTracking()
            .Include(p => p.Maintainers)
            .FirstOrDefaultAsync(p => p.Id == id, cancellationToken);

        if (person == null) return null;

        // R4: Gate person visibility
        bool canSeePerson = accessLevel switch {
            < EAccessLevel.Maintainer => person.Visibility <= EVisibility.Protected,
            EAccessLevel.Maintainer when userId.HasValue => person.Visibility <= EVisibility.Protected
                || await context.PersonMaintainers.AnyAsync(pm => pm.UserId == userId.Value && pm.PersonId == id, cancellationToken),
            EAccessLevel.Maintainer => person.Visibility <= EVisibility.Protected,
            _ => true
        };

        if (!canSeePerson) return null;

        bool isMaintainerOfPerson = accessLevel == EAccessLevel.Maintainer && userId.HasValue
            && await context.PersonMaintainers.AnyAsync(pm => pm.UserId == userId.Value && pm.PersonId == id, cancellationToken);

        IQueryable<Album> albumsQuery = context.Albums
            .AsNoTracking()
            .Where(a => a.PersonOwnerId == id);

        // Album search — ILike so the trgm GIN index on Title can be used
        if (!string.IsNullOrEmpty(albumSearch))
            albumsQuery = albumsQuery.Where(a => EF.Functions.ILike(a.Title, $"%{albumSearch}%"));

        // Album visibility filter. Admin, curators, and maintainers of this person see all albums;
        // everyone else sees Protected-or-below albums plus albums that contain at least one visible asset.
        bool seesAllAlbums = accessLevel >= EAccessLevel.Curator || isMaintainerOfPerson;
        if (!seesAllAlbums) {
            albumsQuery = albumsQuery.Where(a => a.Visibility <= EVisibility.Protected || a.Assets!.Any(
                asset => asset.DeletedAt == null && (
                    asset.Visibility == EVisibility.Public ||
                    (asset.Visibility == EVisibility.Protected && userId.HasValue) ||
                    (asset.Visibility == EVisibility.Private && asset.UploadedBy == userId)
                )));
        }

        IOrderedQueryable<Album> ordered;
        switch (albumSortBy?.ToLowerInvariant()) {
            case "name":
                ordered = albumSortAsc ? albumsQuery.OrderBy(a => a.Title) : albumsQuery.OrderByDescending(a => a.Title);
                break;
            case "created":
                ordered = albumSortAsc ? albumsQuery.OrderBy(a => a.CreatedAt) : albumsQuery.OrderByDescending(a => a.CreatedAt);
                break;
            case "updated":
                ordered = albumSortAsc ? albumsQuery.OrderBy(a => a.UpdatedAt) : albumsQuery.OrderByDescending(a => a.UpdatedAt);
                break;
            case "assets":
                if (seesAllAlbums) {
                    ordered = albumSortAsc
                        ? albumsQuery.OrderBy(a => a.Assets!.Count(asset => asset.DeletedAt == null))
                        : albumsQuery.OrderByDescending(a => a.Assets!.Count(asset => asset.DeletedAt == null));
                } else {
                    ordered = albumSortAsc
                        ? albumsQuery.OrderBy(a => a.Assets!.Count(asset => asset.DeletedAt == null && (
                            asset.Visibility == EVisibility.Public ||
                            (asset.Visibility == EVisibility.Protected && userId.HasValue) ||
                            (asset.Visibility == EVisibility.Private && asset.UploadedBy == userId))))
                        : albumsQuery.OrderByDescending(a => a.Assets!.Count(asset => asset.DeletedAt == null && (
                            asset.Visibility == EVisibility.Public ||
                            (asset.Visibility == EVisibility.Protected && userId.HasValue) ||
                            (asset.Visibility == EVisibility.Private && asset.UploadedBy == userId))));
                }
                break;
            default:
                ordered = albumSortAsc ? albumsQuery.OrderBy(a => a.CreatedAt) : albumsQuery.OrderByDescending(a => a.CreatedAt);
                break;
        }

        // Total album count (after search and visibility filters, before pagination)
        person.AlbumTotalCount = await ordered.CountAsync(cancellationToken);

        // Total visible assets across all visible albums of the person (before search and pagination)
        if (seesAllAlbums) {
            person.TotalAssetCount = await context.Assets
                .Where(a => a.Albums!.Any(al => al.PersonOwnerId == id))
                .CountAsync(asset => asset.DeletedAt == null, cancellationToken);
        } else {
            person.TotalAssetCount = await context.Assets
                .Where(a => a.Albums!.Any(al => al.PersonOwnerId == id))
                .CountAsync(asset => asset.DeletedAt == null && (
                    asset.Visibility == EVisibility.Public ||
                    (asset.Visibility == EVisibility.Protected && userId.HasValue) ||
                    (asset.Visibility == EVisibility.Private && asset.UploadedBy == userId)), cancellationToken);
        }

        var albumIds = await ordered
            .Skip(albumPage * albumPageSize)
            .Take(albumPageSize)
            .Select(a => a.Id)
            .ToListAsync(cancellationToken);

        if (albumIds.Count == 0)
            return person;

        person.Albums = await context.Albums
            .AsNoTracking()
            .AsSplitQuery()
            .Include(a => a.PersonOwner)
            .Include(a => a.CoverAsset)
            .Where(a => albumIds.Contains(a.Id))
            .ToListAsync(cancellationToken);

        // Per-album visible asset counts at the database level
        if (seesAllAlbums) {
            person.AlbumAssetCounts = await context.Albums
                .Where(a => albumIds.Contains(a.Id))
                .Select(a => new { a.Id, Count = a.Assets!.Count(asset => asset.DeletedAt == null) })
                .ToDictionaryAsync(x => x.Id, x => x.Count, cancellationToken);
        } else {
            person.AlbumAssetCounts = await context.Albums
                .Where(a => albumIds.Contains(a.Id))
                .Select(a => new {
                    a.Id,
                    Count = a.Assets!.Count(asset => asset.DeletedAt == null && (
                        asset.Visibility == EVisibility.Public ||
                        (asset.Visibility == EVisibility.Protected && userId.HasValue) ||
                        (asset.Visibility == EVisibility.Private && asset.UploadedBy == userId)))
                })
                .ToDictionaryAsync(x => x.Id, x => x.Count, cancellationToken);
        }

        var orderMap = albumIds.Select((id, i) => (id, i)).ToDictionary(x => x.id, x => x.i);
        person.Albums = [.. person.Albums.OrderBy(a => orderMap.GetValueOrDefault(a.Id))];

        return person;
    }

    /// <inheritdoc />
    public Task<List<Person>> GetAllAsync(CancellationToken cancellationToken) =>
        context.People.Include(p => p.Albums).ToListAsync(cancellationToken);

    /// <inheritdoc />
    public async Task<List<Person>> GetAllVisibleAsync(Guid userId, EAccessLevel accessLevel, CancellationToken cancellationToken) =>
        accessLevel switch {
            // Admin and curator.
            >= EAccessLevel.Curator => await GetAllAsync(cancellationToken),

            EAccessLevel.Maintainer => await context.People
                    .Include(p => p.Albums)
                    .Where(p => p.Visibility <= EVisibility.Protected
                        || context.PersonMaintainers.Any(pm => pm.UserId == userId && pm.PersonId == p.Id))
                    .ToListAsync(cancellationToken),

            EAccessLevel.User => await context.People
                    .Include(p => p.Albums)
                    .Where(p => p.Visibility <= EVisibility.Protected)
                    .ToListAsync(cancellationToken),

            _ => throw new ArgumentOutOfRangeException(nameof(accessLevel), accessLevel, null)
        };

    /// <inheritdoc />
    public Task<Person?> FindByNameAsync(string name, CancellationToken cancellationToken) =>
        context.People.FirstOrDefaultAsync(p => p.Name == name, cancellationToken);

    /// <inheritdoc />
    public Task<List<Person>> FindByNamesAsync(IEnumerable<string> names, CancellationToken cancellationToken) =>
        context.People.Where(p => names.Contains(p.Name)).ToListAsync(cancellationToken);

    /// <inheritdoc />
    public Task<List<Person>> FindBulkAsync(IEnumerable<Guid> ids, CancellationToken cancellationToken) =>
        context.People.Where(p => ids.Contains(p.Id)).ToListAsync(cancellationToken);

    /// <inheritdoc />
    public void Remove(Person person) => context.People.Remove(person);

    /// <inheritdoc />
    public async Task<List<PersonPreviewDto>> SearchQueryAsync(string query, int page, int pageSize, string? sortBy, bool sortAsc, Guid userId, EAccessLevel accessLevel, CancellationToken cancellationToken) {
        IQueryable<Person> peopleQuery = context.People;

        if (!string.IsNullOrEmpty(query))
            peopleQuery = peopleQuery.Where(p => EF.Functions.ILike(p.Name, $"%{query}%"));

        // Apply visibility filter
        peopleQuery = accessLevel switch {
            EAccessLevel.Admin or EAccessLevel.Curator => peopleQuery,

            EAccessLevel.Maintainer => peopleQuery.Where(p => p.Visibility <= EVisibility.Protected
                    || context.PersonMaintainers.Any(pm => pm.UserId == userId && pm.PersonId == p.Id)),

            EAccessLevel.User => peopleQuery.Where(p => p.Visibility <= EVisibility.Protected),

            _  => throw new ArgumentOutOfRangeException(nameof(accessLevel), accessLevel, null)
        };

        // Apply sorting — use subquery counts for "albums"
        IOrderedQueryable<Person> ordered = sortBy?.ToLower() switch
        {
            "created" => sortAsc ? peopleQuery.OrderBy(p => p.CreatedAt) : peopleQuery.OrderByDescending(p => p.CreatedAt),
            "albums" => sortAsc
                ? peopleQuery.OrderBy(p => p.Albums!.Count)
                : peopleQuery.OrderByDescending(p => p.Albums!.Count),
            _ => sortAsc ? peopleQuery.OrderBy(p => p.Name) : peopleQuery.OrderByDescending(p => p.Name),
        };

        var pageOffset = page * pageSize;
        var personIds = await ordered
            .Skip(pageOffset)
            .Take(pageSize)
            .Select(p => p.Id)
            .ToListAsync(cancellationToken);

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

        var pagedPeople = await context.People
            .AsNoTracking()
            .Where(p => personIds.Contains(p.Id))
            .ToListAsync(cancellationToken);

        // Album counts via a single grouped query instead of loading every album row
        var albumCounts = await context.People
            .Where(p => personIds.Contains(p.Id))
            .Select(p => new { p.Id, Count = p.Albums!.Count })
            .ToDictionaryAsync(x => x.Id, x => x.Count, cancellationToken);

        var dtos = pagedPeople.Select(p => {
            var dto = p.ToPersonPreviewDto(accessLevel);
            dto.TotalAlbums = albumCounts.GetValueOrDefault(p.Id, 0);
            return dto;
        }).ToList();

        var orderMap = personIds.Select((id, i) => (id, i)).ToDictionary(x => x.id, x => x.i);
        return [.. dtos.OrderBy(d => orderMap.GetValueOrDefault(d.Id))];
    }

    /// <inheritdoc />
    public Task<bool> IsMaintainerOfAsync(Guid userId, Guid personId, CancellationToken cancellationToken) =>
        context.PersonMaintainers.AnyAsync(pm => pm.UserId == userId && pm.PersonId == personId, cancellationToken);

    /// <inheritdoc />
    public async Task SetMaintainersAsync(Guid personId, IEnumerable<Guid> userIds, CancellationToken cancellationToken) {
        var existing = await context.PersonMaintainers
            .Where(pm => pm.PersonId == personId)
            .ToListAsync(cancellationToken);
        context.PersonMaintainers.RemoveRange(existing);

        foreach (var userId in userIds) {
            context.PersonMaintainers.Add(new PersonMaintainer { PersonId = personId, UserId = userId });
        }
    }

    /// <inheritdoc />
    public async Task MergePeopleAsync(Guid destId, List<Guid> sourceIds, CancellationToken cancellationToken) {
        var dest = await context.People
            .AsSplitQuery()
            .Include(p => p.Maintainers)
            .FirstOrDefaultAsync(p => p.Id == destId, cancellationToken);
        if (dest == null) return;

        // Reassign albums
        var albums = await context.Albums.Where(a => a.PersonOwnerId != null && sourceIds.Contains(a.PersonOwnerId.Value)).ToListAsync(cancellationToken);
        foreach (var album in albums)
            album.PersonOwnerId = destId;

        // Reassign faces
        var faces = await context.Faces.Where(f => f.PersonId.HasValue && sourceIds.Contains(f.PersonId.Value)).ToListAsync(cancellationToken);
        foreach (var face in faces)
            face.PersonId = destId;

        // Reassign maintainers
        var existingMaintainerIds = dest.Maintainers?.Select(m => m.Id).ToHashSet() ?? [];
        var sourceMaintainers = await context.PersonMaintainers
            .Where(pm => sourceIds.Contains(pm.PersonId))
            .ToListAsync(cancellationToken);
        foreach (var pm in sourceMaintainers) {
            if (!existingMaintainerIds.Contains(pm.UserId))
                context.PersonMaintainers.Add(new PersonMaintainer { PersonId = destId, UserId = pm.UserId });
        }

        // Hard-delete source people
        var people = await context.People.Where(p => sourceIds.Contains(p.Id)).ToListAsync(cancellationToken);
        context.People.RemoveRange(people);
    }

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