using System.Threading;
using System.Threading.Tasks;
using MongoDB.Driver;

namespace StellaOps.Vexer.Storage.Mongo.Migrations;

internal sealed class VexInitialIndexMigration : IVexMongoMigration
{
    public string Id => "20251016-initial-indexes";

    public async ValueTask ExecuteAsync(IMongoDatabase database, CancellationToken cancellationToken)
    {
        ArgumentNullException.ThrowIfNull(database);

        await EnsureRawIndexesAsync(database, cancellationToken).ConfigureAwait(false);
        await EnsureProviderIndexesAsync(database, cancellationToken).ConfigureAwait(false);
        await EnsureConsensusIndexesAsync(database, cancellationToken).ConfigureAwait(false);
        await EnsureExportIndexesAsync(database, cancellationToken).ConfigureAwait(false);
        await EnsureCacheIndexesAsync(database, cancellationToken).ConfigureAwait(false);
    }

    private static Task EnsureRawIndexesAsync(IMongoDatabase database, CancellationToken cancellationToken)
    {
        var collection = database.GetCollection<VexRawDocumentRecord>(VexMongoCollectionNames.Raw);
        var providerFormatIndex = Builders<VexRawDocumentRecord>.IndexKeys
            .Ascending(x => x.ProviderId)
            .Ascending(x => x.Format)
            .Ascending(x => x.RetrievedAt);
        return collection.Indexes.CreateOneAsync(new CreateIndexModel<VexRawDocumentRecord>(providerFormatIndex), cancellationToken: cancellationToken);
    }

    private static Task EnsureProviderIndexesAsync(IMongoDatabase database, CancellationToken cancellationToken)
    {
        var collection = database.GetCollection<VexProviderRecord>(VexMongoCollectionNames.Providers);
        var kindIndex = Builders<VexProviderRecord>.IndexKeys.Ascending(x => x.Kind);
        return collection.Indexes.CreateOneAsync(new CreateIndexModel<VexProviderRecord>(kindIndex), cancellationToken: cancellationToken);
    }

    private static Task EnsureConsensusIndexesAsync(IMongoDatabase database, CancellationToken cancellationToken)
    {
        var collection = database.GetCollection<VexConsensusRecord>(VexMongoCollectionNames.Consensus);
        var vulnProductIndex = Builders<VexConsensusRecord>.IndexKeys
            .Ascending(x => x.VulnerabilityId)
            .Ascending(x => x.Product.Key);
        var policyIndex = Builders<VexConsensusRecord>.IndexKeys
            .Ascending(x => x.PolicyRevisionId)
            .Ascending(x => x.PolicyDigest);

        return Task.WhenAll(
            collection.Indexes.CreateOneAsync(new CreateIndexModel<VexConsensusRecord>(vulnProductIndex, new CreateIndexOptions { Unique = true }), cancellationToken: cancellationToken),
            collection.Indexes.CreateOneAsync(new CreateIndexModel<VexConsensusRecord>(policyIndex), cancellationToken: cancellationToken));
    }

    private static Task EnsureExportIndexesAsync(IMongoDatabase database, CancellationToken cancellationToken)
    {
        var collection = database.GetCollection<VexExportManifestRecord>(VexMongoCollectionNames.Exports);
        var signatureIndex = Builders<VexExportManifestRecord>.IndexKeys
            .Ascending(x => x.QuerySignature)
            .Ascending(x => x.Format);
        return collection.Indexes.CreateOneAsync(new CreateIndexModel<VexExportManifestRecord>(signatureIndex, new CreateIndexOptions { Unique = true }), cancellationToken: cancellationToken);
    }

    private static Task EnsureCacheIndexesAsync(IMongoDatabase database, CancellationToken cancellationToken)
    {
        var collection = database.GetCollection<VexCacheEntryRecord>(VexMongoCollectionNames.Cache);
        var signatureIndex = Builders<VexCacheEntryRecord>.IndexKeys
            .Ascending(x => x.QuerySignature)
            .Ascending(x => x.Format);
        var expirationIndex = Builders<VexCacheEntryRecord>.IndexKeys.Ascending(x => x.ExpiresAt);

        return Task.WhenAll(
            collection.Indexes.CreateOneAsync(new CreateIndexModel<VexCacheEntryRecord>(signatureIndex, new CreateIndexOptions { Unique = true }), cancellationToken: cancellationToken),
            collection.Indexes.CreateOneAsync(new CreateIndexModel<VexCacheEntryRecord>(expirationIndex, new CreateIndexOptions { ExpireAfter = TimeSpan.FromSeconds(0) }), cancellationToken: cancellationToken));
    }
}
