CBDD/src/CBDD.Core/Context/DocumentDbContext.cs

using ZB.MOM.WW.CBDD.Bson;
using ZB.MOM.WW.CBDD.Core.CDC;
using ZB.MOM.WW.CBDD.Core.Collections;
using ZB.MOM.WW.CBDD.Core.Compression;
using ZB.MOM.WW.CBDD.Core.Metadata;
using ZB.MOM.WW.CBDD.Core.Storage;
using ZB.MOM.WW.CBDD.Core.Transactions;

namespace ZB.MOM.WW.CBDD.Core;

internal interface ICompactionAwareCollection
{
    /// <summary>
    ///     Refreshes index bindings after compaction.
    /// </summary>
    void RefreshIndexBindingsAfterCompaction();
}

/// <summary>
///     Base class for database contexts.
///     Inherit and add DocumentCollection{T} properties for your entities.
///     Use partial class for Source Generator integration.
/// </summary>
public abstract class DocumentDbContext : IDisposable, ITransactionHolder
{
    internal readonly ChangeStreamDispatcher _cdc;
    private readonly List<ICompactionAwareCollection> _compactionAwareCollections = new();

    private readonly IReadOnlyDictionary<Type, object> _model;
    private readonly List<IDocumentMapper> _registeredMappers = new();
    private readonly IStorageEngine _storage;
    private readonly SemaphoreSlim _transactionLock = new(1, 1);
    protected bool _disposed;

    /// <summary>
    ///     Creates a new database context with default configuration
    /// </summary>
    /// <param name="databasePath">The database file path.</param>
    protected DocumentDbContext(string databasePath)
        : this(databasePath, PageFileConfig.Default, CompressionOptions.Default)
    {
    }

    /// <summary>
    ///     Creates a new database context with default storage configuration and custom compression settings.
    /// </summary>
    /// <param name="databasePath">The database file path.</param>
    /// <param name="compressionOptions">Compression behavior options.</param>
    protected DocumentDbContext(string databasePath, CompressionOptions compressionOptions)
        : this(databasePath, PageFileConfig.Default, compressionOptions)
    {
    }

    /// <summary>
    ///     Creates a new database context with custom configuration
    /// </summary>
    /// <param name="databasePath">The database file path.</param>
    /// <param name="config">The page file configuration.</param>
    protected DocumentDbContext(string databasePath, PageFileConfig config)
        : this(databasePath, config, CompressionOptions.Default)
    {
    }

    /// <summary>
    ///     Creates a new database context with custom storage and compression configuration.
    /// </summary>
    /// <param name="databasePath">The database file path.</param>
    /// <param name="config">The page file configuration.</param>
    /// <param name="compressionOptions">Compression behavior options.</param>
    /// <param name="maintenanceOptions">Maintenance scheduling options.</param>
    protected DocumentDbContext(
        string databasePath,
        PageFileConfig config,
        CompressionOptions? compressionOptions,
        MaintenanceOptions? maintenanceOptions = null)
    {
        if (string.IsNullOrWhiteSpace(databasePath))
            throw new ArgumentNullException(nameof(databasePath));

        _storage = new StorageEngine(databasePath, config, compressionOptions, maintenanceOptions);
        _cdc = new ChangeStreamDispatcher();
        _storage.RegisterCdc(_cdc);

        // Initialize model before collections
        var modelBuilder = new ModelBuilder();
        OnModelCreating(modelBuilder);
        _model = modelBuilder.GetEntityBuilders();
        InitializeCollections();
    }

    /// <summary>
    ///     Gets the current active transaction, if any.
    /// </summary>
    public ITransaction? CurrentTransaction
    {
        get
        {
            if (_disposed)
                throw new ObjectDisposedException(nameof(DocumentDbContext));
            return field != null && field.State == TransactionState.Active ? field : null;
        }
        private set;
    }

    /// <summary>
    ///     Gets the concrete storage engine for advanced scenarios in derived contexts.
    /// </summary>
    protected StorageEngine Engine => (StorageEngine)_storage;

    /// <summary>
    ///     Gets compression options bound to this context's storage engine.
    /// </summary>
    protected CompressionOptions CompressionOptions => _storage.CompressionOptions;

    /// <summary>
    ///     Gets the compression service for codec operations.
    /// </summary>
    protected CompressionService CompressionService => _storage.CompressionService;

    /// <summary>
    ///     Gets compression telemetry counters.
    /// </summary>
    protected CompressionTelemetry CompressionTelemetry => _storage.CompressionTelemetry;

    /// <summary>
    ///     Releases resources used by the context.
    /// </summary>
    public void Dispose()
    {
        if (_disposed)
            return;

        _disposed = true;

        _storage?.Dispose();
        _cdc?.Dispose();
        _transactionLock?.Dispose();

        GC.SuppressFinalize(this);
    }

    /// <summary>
    ///     Gets the current active transaction or starts a new one.
    /// </summary>
    /// <returns>The active transaction.</returns>
    public ITransaction GetCurrentTransactionOrStart()
    {
        return BeginTransaction();
    }

    /// <summary>
    ///     Gets the current active transaction or starts a new one asynchronously.
    /// </summary>
    /// <returns>The active transaction.</returns>
    public async Task<ITransaction> GetCurrentTransactionOrStartAsync()
    {
        return await BeginTransactionAsync();
    }

    /// <summary>
    ///     Initializes document collections for the context.
    /// </summary>
    protected virtual void InitializeCollections()
    {
        // Derived classes can override to initialize collections
    }

    /// <summary>
    ///     Override to configure the model using Fluent API.
    /// </summary>
    /// <param name="modelBuilder">The model builder instance.</param>
    protected virtual void OnModelCreating(ModelBuilder modelBuilder)
    {
    }

    /// <summary>
    ///     Helper to create a DocumentCollection instance with custom TId.
    ///     Used by derived classes in InitializeCollections for typed primary keys.
    /// </summary>
    /// <typeparam name="TId">The document identifier type.</typeparam>
    /// <typeparam name="T">The document type.</typeparam>
    /// <param name="mapper">The mapper used for document serialization and key access.</param>
    /// <returns>The created document collection.</returns>
    protected DocumentCollection<TId, T> CreateCollection<TId, T>(IDocumentMapper<TId, T> mapper)
        where T : class
    {
        if (_disposed)
            throw new ObjectDisposedException(nameof(DocumentDbContext));

        string? customName = null;
        EntityTypeBuilder<T>? builder = null;

        if (_model.TryGetValue(typeof(T), out object? builderObj))
        {
            builder = builderObj as EntityTypeBuilder<T>;
            customName = builder?.CollectionName;
        }

        _registeredMappers.Add(mapper);
        var collection = new DocumentCollection<TId, T>(_storage, this, mapper, customName);
        if (collection is ICompactionAwareCollection compactionAwareCollection)
            _compactionAwareCollections.Add(compactionAwareCollection);

        // Apply configurations from ModelBuilder
        if (builder != null)
            foreach (var indexBuilder in builder.Indexes)
                collection.ApplyIndexBuilder(indexBuilder);

        _storage.RegisterMappers(_registeredMappers);

        return collection;
    }

    /// <summary>
    ///     Gets the document collection for the specified entity type using an ObjectId as the key.
    /// </summary>
    /// <typeparam name="T">The type of entity to retrieve the document collection for. Must be a reference type.</typeparam>
    /// <returns>A DocumentCollection&lt;ObjectId, T&gt; instance for the specified entity type.</returns>
    public DocumentCollection<ObjectId, T> Set<T>() where T : class
    {
        return Set<ObjectId, T>();
    }

    /// <summary>
    ///     Gets a collection for managing documents of type T, identified by keys of type TId.
    ///     Override is generated automatically by the Source Generator for partial DbContext classes.
    /// </summary>
    /// <typeparam name="TId">The type of the unique identifier for documents in the collection.</typeparam>
    /// <typeparam name="T">The type of the document to be managed. Must be a reference type.</typeparam>
    /// <returns>A DocumentCollection&lt;TId, T&gt; instance for performing operations on documents of type T.</returns>
    public virtual DocumentCollection<TId, T> Set<TId, T>() where T : class
    {
        throw new InvalidOperationException(
            $"No collection registered for entity type '{typeof(T).Name}' with key type '{typeof(TId).Name}'.");
    }

    /// <summary>
    ///     Begins a transaction or returns the current active transaction.
    /// </summary>
    /// <returns>The active transaction.</returns>
    public ITransaction BeginTransaction()
    {
        if (_disposed)
            throw new ObjectDisposedException(nameof(DocumentDbContext));

        _transactionLock.Wait();
        try
        {
            if (CurrentTransaction != null)
                return CurrentTransaction; // Return existing active transaction
            CurrentTransaction = _storage.BeginTransaction();
            return CurrentTransaction;
        }
        finally
        {
            _transactionLock.Release();
        }
    }

    /// <summary>
    ///     Begins a transaction asynchronously or returns the current active transaction.
    /// </summary>
    /// <param name="ct">The cancellation token.</param>
    /// <returns>The active transaction.</returns>
    public async Task<ITransaction> BeginTransactionAsync(CancellationToken ct = default)
    {
        if (_disposed)
            throw new ObjectDisposedException(nameof(DocumentDbContext));

        var lockAcquired = false;
        try
        {
            await _transactionLock.WaitAsync(ct);
            lockAcquired = true;

            if (CurrentTransaction != null)
                return CurrentTransaction; // Return existing active transaction
            CurrentTransaction = await _storage.BeginTransactionAsync(IsolationLevel.ReadCommitted, ct);
            return CurrentTransaction;
        }
        finally
        {
            if (lockAcquired)
                _transactionLock.Release();
        }
    }

    /// <summary>
    ///     Commits the current transaction if one is active.
    /// </summary>
    public void SaveChanges()
    {
        if (_disposed)
            throw new ObjectDisposedException(nameof(DocumentDbContext));
        if (CurrentTransaction != null)
            try
            {
                CurrentTransaction.Commit();
            }
            finally
            {
                CurrentTransaction = null;
            }
    }

    /// <summary>
    ///     Commits the current transaction asynchronously if one is active.
    /// </summary>
    /// <param name="ct">The cancellation token.</param>
    public async Task SaveChangesAsync(CancellationToken ct = default)
    {
        if (_disposed)
            throw new ObjectDisposedException(nameof(DocumentDbContext));
        if (CurrentTransaction != null)
            try
            {
                await CurrentTransaction.CommitAsync(ct);
            }
            finally
            {
                CurrentTransaction = null;
            }
    }

    /// <summary>
    ///     Executes a checkpoint using the requested mode.
    /// </summary>
    /// <param name="mode">Checkpoint mode to execute.</param>
    /// <returns>The checkpoint execution result.</returns>
    public CheckpointResult Checkpoint(CheckpointMode mode = CheckpointMode.Truncate)
    {
        if (_disposed)
            throw new ObjectDisposedException(nameof(DocumentDbContext));

        return Engine.Checkpoint(mode);
    }

    /// <summary>
    ///     Executes a checkpoint asynchronously using the requested mode.
    /// </summary>
    /// <param name="mode">Checkpoint mode to execute.</param>
    /// <param name="ct">The cancellation token.</param>
    /// <returns>The checkpoint execution result.</returns>
    public Task<CheckpointResult> CheckpointAsync(CheckpointMode mode = CheckpointMode.Truncate,
        CancellationToken ct = default)
    {
        if (_disposed)
            throw new ObjectDisposedException(nameof(DocumentDbContext));

        return Engine.CheckpointAsync(mode, ct);
    }

    /// <summary>
    ///     Returns a point-in-time snapshot of compression telemetry counters.
    /// </summary>
    public CompressionStats GetCompressionStats()
    {
        if (_disposed)
            throw new ObjectDisposedException(nameof(DocumentDbContext));

        return Engine.GetCompressionStats();
    }

    /// <summary>
    ///     Runs offline compaction by default. Set options to online mode for a bounded online pass.
    /// </summary>
    /// <param name="options">Compaction execution options.</param>
    public CompactionStats Compact(CompactionOptions? options = null)
    {
        if (_disposed)
            throw new ObjectDisposedException(nameof(DocumentDbContext));

        var stats = Engine.Compact(options);
        RefreshCollectionBindingsAfterCompaction();
        return stats;
    }

    /// <summary>
    ///     Runs offline compaction by default. Set options to online mode for a bounded online pass.
    /// </summary>
    /// <param name="options">Compaction execution options.</param>
    /// <param name="ct">Cancellation token for the asynchronous operation.</param>
    public Task<CompactionStats> CompactAsync(CompactionOptions? options = null, CancellationToken ct = default)
    {
        if (_disposed)
            throw new ObjectDisposedException(nameof(DocumentDbContext));

        return CompactAsyncCore(options, ct);
    }

    /// <summary>
    ///     Alias for <see cref="Compact(CompactionOptions?)" />.
    /// </summary>
    /// <param name="options">Compaction execution options.</param>
    public CompactionStats Vacuum(CompactionOptions? options = null)
    {
        if (_disposed)
            throw new ObjectDisposedException(nameof(DocumentDbContext));

        var stats = Engine.Vacuum(options);
        RefreshCollectionBindingsAfterCompaction();
        return stats;
    }

    /// <summary>
    ///     Async alias for <see cref="CompactAsync(CompactionOptions?, CancellationToken)" />.
    /// </summary>
    /// <param name="options">Compaction execution options.</param>
    /// <param name="ct">Cancellation token for the asynchronous operation.</param>
    public Task<CompactionStats> VacuumAsync(CompactionOptions? options = null, CancellationToken ct = default)
    {
        if (_disposed)
            throw new ObjectDisposedException(nameof(DocumentDbContext));

        return VacuumAsyncCore(options, ct);
    }

    private async Task<CompactionStats> CompactAsyncCore(CompactionOptions? options, CancellationToken ct)
    {
        var stats = await Engine.CompactAsync(options, ct);
        RefreshCollectionBindingsAfterCompaction();
        return stats;
    }

    private async Task<CompactionStats> VacuumAsyncCore(CompactionOptions? options, CancellationToken ct)
    {
        var stats = await Engine.VacuumAsync(options, ct);
        RefreshCollectionBindingsAfterCompaction();
        return stats;
    }

    private void RefreshCollectionBindingsAfterCompaction()
    {
        foreach (var collection in _compactionAwareCollections) collection.RefreshIndexBindingsAfterCompaction();
    }

    /// <summary>
    ///     Gets page usage grouped by page type.
    /// </summary>
    public IReadOnlyList<PageTypeUsageEntry> GetPageUsageByPageType()
    {
        if (_disposed)
            throw new ObjectDisposedException(nameof(DocumentDbContext));

        return Engine.GetPageUsageByPageType();
    }

    /// <summary>
    ///     Gets per-collection page usage diagnostics.
    /// </summary>
    public IReadOnlyList<CollectionPageUsageEntry> GetPageUsageByCollection()
    {
        if (_disposed)
            throw new ObjectDisposedException(nameof(DocumentDbContext));

        return Engine.GetPageUsageByCollection();
    }

    /// <summary>
    ///     Gets per-collection compression ratio diagnostics.
    /// </summary>
    public IReadOnlyList<CollectionCompressionRatioEntry> GetCompressionRatioByCollection()
    {
        if (_disposed)
            throw new ObjectDisposedException(nameof(DocumentDbContext));

        return Engine.GetCompressionRatioByCollection();
    }

    /// <summary>
    ///     Gets free-list summary diagnostics.
    /// </summary>
    public FreeListSummary GetFreeListSummary()
    {
        if (_disposed)
            throw new ObjectDisposedException(nameof(DocumentDbContext));

        return Engine.GetFreeListSummary();
    }

    /// <summary>
    ///     Gets page-level fragmentation diagnostics.
    /// </summary>
    public FragmentationMapReport GetFragmentationMap()
    {
        if (_disposed)
            throw new ObjectDisposedException(nameof(DocumentDbContext));

        return Engine.GetFragmentationMap();
    }

    /// <summary>
    ///     Runs compression migration as dry-run estimation by default.
    /// </summary>
    /// <param name="options">Compression migration options.</param>
    public CompressionMigrationResult MigrateCompression(CompressionMigrationOptions? options = null)
    {
        if (_disposed)
            throw new ObjectDisposedException(nameof(DocumentDbContext));

        return Engine.MigrateCompression(options);
    }

    /// <summary>
    ///     Runs compression migration asynchronously as dry-run estimation by default.
    /// </summary>
    /// <param name="options">Compression migration options.</param>
    /// <param name="ct">Cancellation token for the asynchronous operation.</param>
    public Task<CompressionMigrationResult> MigrateCompressionAsync(CompressionMigrationOptions? options = null,
        CancellationToken ct = default)
    {
        if (_disposed)
            throw new ObjectDisposedException(nameof(DocumentDbContext));

        return Engine.MigrateCompressionAsync(options, ct);
    }
}