using MongoDB.Driver;
using StellaOps.Scanner.Storage.Catalog;
using StellaOps.Scanner.Storage.Mongo;

namespace StellaOps.Scanner.Storage.Repositories;

public sealed class ArtifactRepository
{
    private readonly MongoCollectionProvider _collections;

    public ArtifactRepository(MongoCollectionProvider collections)
    {
        _collections = collections ?? throw new ArgumentNullException(nameof(collections));
    }

    public async Task<ArtifactDocument?> GetAsync(string artifactId, CancellationToken cancellationToken)
    {
        ArgumentException.ThrowIfNullOrWhiteSpace(artifactId);
        return await _collections.Artifacts
            .Find(x => x.Id == artifactId)
            .FirstOrDefaultAsync(cancellationToken)
            .ConfigureAwait(false);
    }

    public async Task UpsertAsync(ArtifactDocument document, CancellationToken cancellationToken)
    {
        ArgumentNullException.ThrowIfNull(document);
        document.UpdatedAtUtc = DateTime.UtcNow;
        var options = new ReplaceOptions { IsUpsert = true };
        await _collections.Artifacts
            .ReplaceOneAsync(x => x.Id == document.Id, document, options, cancellationToken)
            .ConfigureAwait(false);
    }

    public async Task UpdateRekorAsync(string artifactId, RekorReference reference, CancellationToken cancellationToken)
    {
        ArgumentException.ThrowIfNullOrWhiteSpace(artifactId);
        ArgumentNullException.ThrowIfNull(reference);

        var update = Builders<ArtifactDocument>.Update
            .Set(x => x.Rekor, reference)
            .Set(x => x.UpdatedAtUtc, DateTime.UtcNow);

        await _collections.Artifacts.UpdateOneAsync(x => x.Id == artifactId, update, cancellationToken: cancellationToken).ConfigureAwait(false);
    }

    public async Task<long> IncrementRefCountAsync(string artifactId, long delta, CancellationToken cancellationToken)
    {
        ArgumentException.ThrowIfNullOrWhiteSpace(artifactId);

        var update = Builders<ArtifactDocument>.Update
            .Inc(x => x.RefCount, delta)
            .Set(x => x.UpdatedAtUtc, DateTime.UtcNow);

        var options = new FindOneAndUpdateOptions<ArtifactDocument, ArtifactDocument>
        {
            ReturnDocument = ReturnDocument.After,
            IsUpsert = false,
        };

        var result = await _collections.Artifacts
            .FindOneAndUpdateAsync<ArtifactDocument, ArtifactDocument>(x => x.Id == artifactId, update, options, cancellationToken)
            .ConfigureAwait(false);

        return result?.RefCount ?? 0;
    }
}
