using System.Diagnostics; using LiteDB.Async; using Microsoft.Extensions.Options; using Octokit; 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 ISettingsManager settingsManager; private readonly DebugOptions debugOptions; private LiteDatabaseAsync? database; public LiteDatabaseAsync Database { get { if (database is null) throw new InvalidOperationException("Database not set before access"); return database; } private set { database?.Dispose(); database = value; } } // Notification events public event EventHandler? CivitModelsChanged; // Collections (Tables) public ILiteCollectionAsync CivitModels => Database.GetCollection("CivitModels"); public ILiteCollectionAsync CivitModelVersions => Database.GetCollection("CivitModelVersions"); public ILiteCollectionAsync CivitModelQueryCache => Database.GetCollection("CivitModelQueryCache"); public ILiteCollectionAsync GithubCache => Database.GetCollection("GithubCache"); public LiteDbContext(ISettingsManager settingsManager, IOptions debugOptions) { this.settingsManager = settingsManager; this.debugOptions = debugOptions.Value; var connectionString = debugOptions.Value.TempDatabase ? ":temp:" : $"Filename={Path.Combine(settingsManager.LibraryDir, "StabilityMatrix.db")};Mode=Shared"; Database = new LiteDatabaseAsync(connectionString); // Register reference fields LiteDBExtensions.Register(m => m.ModelVersions, "CivitModelVersions"); LiteDBExtensions.Register(e => e.Items, "CivitModels"); } 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(); 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(); return (model, version); } public async Task UpsertCivitModelAsync(CivitModel civitModel) { // Insert model versions first then model var versionsUpdated = await CivitModelVersions.UpsertAsync(civitModel.ModelVersions); var updated = await CivitModels.UpsertAsync(civitModel); // Notify listeners on any change var anyUpdated = versionsUpdated > 0 || updated; if (anyUpdated) { CivitModelsChanged?.Invoke(this, EventArgs.Empty); } return anyUpdated; } public async Task UpsertCivitModelAsync(IEnumerable 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); var updated = await CivitModels.UpsertAsync(civitModelsArray); // 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 UpsertCivitModelQueryCacheEntryAsync(CivitModelQueryCacheEntry entry) { var changed = await CivitModelQueryCache.UpsertAsync(entry); if (changed) { CivitModelsChanged?.Invoke(this, EventArgs.Empty); } return changed; } public Task GetGithubCacheEntry(string cacheKey) => GithubCache.FindByIdAsync(cacheKey); public Task UpsertGithubCacheEntry(GithubCacheEntry cacheEntry) => GithubCache.UpsertAsync(cacheEntry); public void Dispose() { database?.UnderlyingDatabase.Dispose(); database?.Dispose(); GC.SuppressFinalize(this); } }