StabilityMatrix/StabilityMatrix.Core/Helper/ModelFinder.cs

using NLog;
using Refit;
using StabilityMatrix.Core.Api;
using StabilityMatrix.Core.Attributes;
using StabilityMatrix.Core.Database;
using StabilityMatrix.Core.Models.Api;

namespace StabilityMatrix.Core.Helper;

// return Model, ModelVersion, ModelFile
public record struct ModelSearchResult(
    CivitModel Model,
    CivitModelVersion ModelVersion,
    CivitFile ModelFile
);

[Singleton]
public class ModelFinder
{
    private static readonly Logger Logger = LogManager.GetCurrentClassLogger();
    private readonly ILiteDbContext liteDbContext;
    private readonly ICivitApi civitApi;

    public ModelFinder(ILiteDbContext liteDbContext, ICivitApi civitApi)
    {
        this.liteDbContext = liteDbContext;
        this.civitApi = civitApi;
    }

    // Finds a model from the local database using file hash
    public async Task<ModelSearchResult?> LocalFindModel(string hashBlake3)
    {
        var (model, version) = await liteDbContext.FindCivitModelFromFileHashAsync(hashBlake3);
        if (model == null || version == null)
        {
            return null;
        }

        var file = version.Files!.First(
            file => file.Hashes.BLAKE3?.ToLowerInvariant() == hashBlake3
        );

        return new ModelSearchResult(model, version, file);
    }

    // Finds a model using Civit API using file hash
    public async Task<ModelSearchResult?> RemoteFindModel(string hashBlake3)
    {
        Logger.Info("Searching Civit API for model version using hash {Hash}", hashBlake3);
        try
        {
            var versionResponse = await civitApi.GetModelVersionByHash(hashBlake3);

            Logger.Info(
                "Found version {VersionId} with model id {ModelId}",
                versionResponse.Id,
                versionResponse.ModelId
            );
            var model = await civitApi.GetModelById(versionResponse.ModelId);

            // VersionResponse is not actually the full data of ModelVersion, so find it again
            var version = model.ModelVersions!.First(version => version.Id == versionResponse.Id);

            var file = versionResponse.Files.First(
                file => file.Hashes.BLAKE3?.ToLowerInvariant() == hashBlake3
            );

            return new ModelSearchResult(model, version, file);
        }
        catch (ApiException e)
        {
            Logger.Info(
                "Could not find remote model version using hash {Hash}: {Error}",
                hashBlake3,
                e.Message
            );
            return null;
        }
    }
}