StabilityMatrix/StabilityMatrix.Core/Database/LiteDbContext.cs

using LiteDB;
using LiteDB.Async;
using Microsoft.Extensions.Logging;
using Microsoft.Extensions.Options;
using StabilityMatrix.Core.Extensions;
using StabilityMatrix.Core.Models.Api;
using StabilityMatrix.Core.Models.Configs;
using StabilityMatrix.Core.Models.Database;
using StabilityMatrix.Core.Services;

namespace StabilityMatrix.Core.Database;

public class LiteDbContext : ILiteDbContext
{
    private readonly ILogger<LiteDbContext> logger;
    private readonly ISettingsManager settingsManager;
    private readonly DebugOptions debugOptions;

    private LiteDatabaseAsync? database;
    public LiteDatabaseAsync Database => database ??= CreateDatabase();

    // Notification events
    public event EventHandler? CivitModelsChanged;
    
    // Collections (Tables)
    public ILiteCollectionAsync<CivitModel> CivitModels => Database.GetCollection<CivitModel>("CivitModels");
    public ILiteCollectionAsync<CivitModelVersion> CivitModelVersions => Database.GetCollection<CivitModelVersion>("CivitModelVersions");
    public ILiteCollectionAsync<CivitModelQueryCacheEntry> CivitModelQueryCache => Database.GetCollection<CivitModelQueryCacheEntry>("CivitModelQueryCache");
    public ILiteCollectionAsync<GithubCacheEntry> GithubCache => Database.GetCollection<GithubCacheEntry>("GithubCache");
    public ILiteCollectionAsync<LocalModelFile> LocalModelFiles => Database.GetCollection<LocalModelFile>("LocalModelFiles");
    
    public LiteDbContext(
        ILogger<LiteDbContext> logger,
        ISettingsManager settingsManager, 
        IOptions<DebugOptions> debugOptions)
    {
        this.logger = logger;
        this.settingsManager = settingsManager;
        this.debugOptions = debugOptions.Value;
    }

    private LiteDatabaseAsync CreateDatabase()
    {
        LiteDatabaseAsync? db = null;
        
        if (debugOptions.TempDatabase)
        {
            db = new LiteDatabaseAsync(":temp:");
        }
        else
        {
            // Attempt to create connection, might be in use
            try
            {
                var dbPath = Path.Combine(settingsManager.LibraryDir, "StabilityMatrix.db");
                db = new LiteDatabaseAsync(new ConnectionString()
                {
                    Filename = dbPath,
                    Connection = ConnectionType.Shared,
                });
            }
            catch (IOException e)
            {
                logger.LogWarning("Database in use or not accessible ({Message}), using temporary database", e.Message);
            }
        }
        
        // Fallback to temporary database
        db ??= new LiteDatabaseAsync(":temp:");

        // Register reference fields
        LiteDBExtensions.Register<CivitModel, CivitModelVersion>(m => m.ModelVersions, "CivitModelVersions");
        LiteDBExtensions.Register<CivitModelQueryCacheEntry, CivitModel>(e => e.Items, "CivitModels");

        return db;
    }
    
    public async Task<(CivitModel?, CivitModelVersion?)> FindCivitModelFromFileHashAsync(string hashBlake3)
    {
        var version = await CivitModelVersions.Query()
            .Where(mv => mv.Files!
                .Select(f => f.Hashes)
                .Select(hashes => hashes.BLAKE3)
                .Any(hash => hash == hashBlake3))
            .FirstOrDefaultAsync()
            .ConfigureAwait(false);
        
        if (version is null) return (null, null);
        
        var model = await CivitModels.Query()
            .Include(m => m.ModelVersions)
            .Where(m => m.ModelVersions!
                .Select(v => v.Id)
                .Any(id => id == version.Id))
            .FirstOrDefaultAsync().ConfigureAwait(false);
        
        return (model, version);
    }
    
    public async Task<bool> UpsertCivitModelAsync(CivitModel civitModel)
    {
        // Insert model versions first then model
        var versionsUpdated = await CivitModelVersions.UpsertAsync(civitModel.ModelVersions).ConfigureAwait(false);
        var updated = await CivitModels.UpsertAsync(civitModel).ConfigureAwait(false);
        // Notify listeners on any change
        var anyUpdated = versionsUpdated > 0 || updated;
        if (anyUpdated)
        {
            CivitModelsChanged?.Invoke(this, EventArgs.Empty);
        }
        return anyUpdated;
    }
    
    public async Task<bool> UpsertCivitModelAsync(IEnumerable<CivitModel> civitModels)
    {
        var civitModelsArray = civitModels.ToArray();
        // Get all model versions then insert models
        var versions = civitModelsArray.SelectMany(model => model.ModelVersions ?? new());
        var versionsUpdated = await CivitModelVersions.UpsertAsync(versions).ConfigureAwait(false);
        var updated = await CivitModels.UpsertAsync(civitModelsArray).ConfigureAwait(false);
        // Notify listeners on any change
        var anyUpdated = versionsUpdated > 0 || updated > 0;
        if (updated > 0 || versionsUpdated > 0)
        {
            CivitModelsChanged?.Invoke(this, EventArgs.Empty);
        }
        return anyUpdated;
    }
    
    // Add to cache
    public async Task<bool> UpsertCivitModelQueryCacheEntryAsync(CivitModelQueryCacheEntry entry)
    {
        var changed = await CivitModelQueryCache.UpsertAsync(entry).ConfigureAwait(false);
        if (changed)
        {
            CivitModelsChanged?.Invoke(this, EventArgs.Empty);
        }

        return changed;
    }

    public async Task<GithubCacheEntry?> GetGithubCacheEntry(string? cacheKey)
    {
        if (string.IsNullOrEmpty(cacheKey)) return null;
        
        if (await GithubCache.FindByIdAsync(cacheKey).ConfigureAwait(false) is { } result)
        {
            return result;
        }
        
        return null;
    }

    public Task<bool> UpsertGithubCacheEntry(GithubCacheEntry cacheEntry) =>
        GithubCache.UpsertAsync(cacheEntry);

    public void Dispose()
    {
        if (database is not null)
        {
            try
            {
                database.Dispose();
            }
            catch (ObjectDisposedException)
            {
            }

            database = null;
        }

        GC.SuppressFinalize(this);
    }
}