Core.Blueprint.DAL/Core.Blueprint.DAL.Infrastrcture/Context/MongoContext.cs

using Azure.Core;
using Azure.Identity;
using Core.Blueprint.DAL.Infrastructure.Contracts;
using Core.Blueprint.Domain.Shared;
using Microsoft.Extensions.Configuration;
using MongoDB.Driver;
using MongoDB.Driver.Authentication.Oidc;
using System.Reflection;

namespace Core.Blueprint.DAL.Infrastructure.Context
{
    public interface IMongoDbSettings
    {
        string DatabaseName { get; set; }
        string ConnectionString { get; set; }
    }

    public class MongoDbSettings : IMongoDbSettings
    {
        public string DatabaseName { get; set; }
        public string ConnectionString { get; set; }

    }

    public class HeathOidcCallback : IOidcCallback
    {
        private readonly string _audience;
        public HeathOidcCallback(string audience)
        {
            _audience = audience;
        }

        public OidcAccessToken GetOidcAccessToken(OidcCallbackParameters parameters, CancellationToken cancellationToken)
        {
            TokenRequestContext tokenRequestContext =
                new TokenRequestContext(
                    new[] { _audience }
                );

            AccessToken token =
                new ChainedTokenCredential(
                    new ManagedIdentityCredential(),
                    new SharedTokenCacheCredential(),
                    new VisualStudioCredential(),
                    new VisualStudioCodeCredential()
                )
                .GetToken(
                    tokenRequestContext
                );

            return new(token.Token, expiresIn: null);
        }

        public async Task<OidcAccessToken> GetOidcAccessTokenAsync(OidcCallbackParameters parameters, CancellationToken cancellationToken)
        {
            TokenRequestContext tokenRequestContext =
                new TokenRequestContext(
                    new[] { _audience }
                );

            AccessToken token =
                await new ChainedTokenCredential(
                    new ManagedIdentityCredential(),
                    new SharedTokenCacheCredential(),
                    new VisualStudioCredential(),
                    new VisualStudioCodeCredential()
                )
                .GetTokenAsync(
                    tokenRequestContext, cancellationToken
                ).ConfigureAwait(false);

            return new(token.Token, expiresIn: null);
        }
    }
    public sealed class MongoContext : IMongoContext
    {
        private IMongoDatabase _database;
        public IMongoClient _client { get; private set; }
        public IMongoDatabase Database { get { return _database; } }
        public MongoUrl _mongoUrl { get; private set; }

        //Constructors
        public MongoContext(IConfiguration configuration)
        {
            var mongoClientSettings = MongoClientSettings.FromConnectionString(configuration.GetConnectionString("MongoDBConnString"));
            mongoClientSettings.Credential = MongoCredential.CreateOidcCredential(new HeathOidcCallback(configuration.GetValue<string>("MongoDb:Audience")));
            _database = new MongoClient(mongoClientSettings).GetDatabase(configuration.GetValue<string>("MongoDb:DatabaseName")); //Mongo Database
            _client = Database.Client;
        }
        public MongoContext(string connectionString, string audience, string databaseName)
        {
            var mongoClientSettings = MongoClientSettings.FromConnectionString(connectionString);
            mongoClientSettings.Credential = MongoCredential.CreateOidcCredential(new HeathOidcCallback(audience));
            _client = new MongoClient(mongoClientSettings);
            _database = _client.GetDatabase(databaseName);
        }
        public MongoContext(MongoClient client, string databaseName)
        {
            _client = client;
            _database = client.GetDatabase(databaseName);
        }
        public void SetUpDatabase(string database)
        {
            _database = new MongoClient(_mongoUrl).GetDatabase(database);
            _client = Database.Client;
        }


        //Methods
        public void DropCollection<TDocument>(string? partitionKey = null) where TDocument : class
        {
            Database.DropCollection(GetCollectionName<TDocument>(partitionKey));
        }
        public IMongoCollection<T> GetCollection<T>(string collection)
        {
            return _database.GetCollection<T>(collection);
        }

        private string GetCollectionName<TDocument>(string partitionKey)
        {
            var collectionName = GetAttributeCollectionName<TDocument>();
            if (string.IsNullOrEmpty(partitionKey))
            {
                return collectionName;
            }
            return $"{partitionKey}-{collectionName}";
        }
        private string? GetAttributeCollectionName<TDocument>()
        {
            return (typeof(TDocument).GetTypeInfo()
                .GetCustomAttributes(typeof(CollectionAttributeName))
                .FirstOrDefault() as CollectionAttributeName)?.Name;
        }
    }
}