using System;
using System.Collections.Generic;
using System.ComponentModel.DataAnnotations;
using System.Diagnostics.CodeAnalysis;
using System.Linq;
using System.Runtime.CompilerServices;
using System.Text.Json.Serialization;
using System.Threading.Tasks;
using AsyncAwaitBestPractices;
using Avalonia.Input;
using Avalonia.Platform.Storage;
using Avalonia.Threading;
using CommunityToolkit.Mvvm.ComponentModel;
using CommunityToolkit.Mvvm.Input;
using NLog;
using StabilityMatrix.Avalonia.Controls;
using StabilityMatrix.Avalonia.Extensions;
using StabilityMatrix.Avalonia.Models;
using StabilityMatrix.Avalonia.Models.Inference;
using StabilityMatrix.Avalonia.Services;
using StabilityMatrix.Avalonia.ViewModels.Base;
using StabilityMatrix.Core.Attributes;
using StabilityMatrix.Core.Models.Database;
using Size = System.Drawing.Size;
#pragma warning disable CS0657 // Not a valid attribute location for this declaration
namespace StabilityMatrix.Avalonia.ViewModels.Inference;
[View(typeof(SelectImageCard))]
[ManagedService]
[Transient]
public partial class SelectImageCardViewModel(INotificationService notificationService)
: LoadableViewModelBase,
IDropTarget,
IComfyStep,
IInputImageProvider
{
private static readonly Logger Logger = LogManager.GetCurrentClassLogger();
[ObservableProperty]
[NotifyPropertyChangedFor(nameof(IsSelectionAvailable))]
[NotifyPropertyChangedFor(nameof(IsImageFileNotFound))]
private ImageSource? imageSource;
[ObservableProperty]
[property: JsonIgnore]
[NotifyPropertyChangedFor(nameof(IsSelectionAvailable))]
private bool isSelectionEnabled = true;
///
/// Set by when the image is loaded.
///
[ObservableProperty]
private Size currentBitmapSize = Size.Empty;
///
/// True if the image file is set but the local file does not exist.
///
[MemberNotNullWhen(true, nameof(NotFoundImagePath))]
public bool IsImageFileNotFound => ImageSource?.LocalFile?.Exists == false;
public bool IsSelectionAvailable => IsSelectionEnabled && ImageSource == null;
///
/// Path of the not found image
///
public string? NotFoundImagePath => ImageSource?.LocalFile?.FullPath;
///
public void ApplyStep(ModuleApplyStepEventArgs e)
{
e.Builder.SetupImagePrimarySource(
ImageSource ?? throw new ValidationException("Input Image is required"),
!CurrentBitmapSize.IsEmpty
? CurrentBitmapSize
: throw new ValidationException("CurrentBitmapSize is null"),
e.Builder.Connections.BatchIndex
);
}
///
public IEnumerable GetInputImages()
{
if (ImageSource is { } image && !IsImageFileNotFound)
{
yield return image;
}
}
partial void OnImageSourceChanged(ImageSource? value)
{
// Cache the hash for later upload use
if (value?.LocalFile is { Exists: true })
{
value.GetBlake3HashAsync().SafeFireAndForget();
}
}
private static FilePickerFileType SupportedImages { get; } =
new("Supported Images")
{
Patterns = new[] { "*.png", "*.jpg", "*.jpeg" },
AppleUniformTypeIdentifiers = new[] { "public.jpeg", "public.png" },
MimeTypes = new[] { "image/jpeg", "image/png" }
};
[RelayCommand]
private async Task SelectImageFromFilePickerAsync()
{
var files = await App.StorageProvider.OpenFilePickerAsync(
new FilePickerOpenOptions
{
FileTypeFilter = [FilePickerFileTypes.ImagePng, FilePickerFileTypes.ImageJpg, SupportedImages]
}
);
if (files.FirstOrDefault()?.TryGetLocalPath() is { } path)
{
Dispatcher.UIThread.Post(() => LoadUserImageSafe(new ImageSource(path)));
}
}
///
/// Supports LocalImageFile Context or OS Files
///
public void DragOver(object? sender, DragEventArgs e)
{
if (
e.Data.GetDataFormats().Contains(DataFormats.Files)
|| e.Data.GetContext() is not null
)
{
e.Handled = true;
return;
}
e.DragEffects = DragDropEffects.None;
}
///
public void Drop(object? sender, DragEventArgs e)
{
// 1. Context drop for LocalImageFile
if (e.Data.GetContext() is { } imageFile)
{
e.Handled = true;
Dispatcher.UIThread.Post(() => LoadUserImageSafe(new ImageSource(imageFile.AbsolutePath)));
return;
}
// 2. OS Files
if (
e.Data.GetFiles() is { } files
&& files.Select(f => f.TryGetLocalPath()).FirstOrDefault() is { } path
)
{
e.Handled = true;
Dispatcher.UIThread.Post(() => LoadUserImageSafe(new ImageSource(path)));
}
}
///
/// Calls with notification error handling.
///
private void LoadUserImageSafe(ImageSource image)
{
try
{
LoadUserImage(image);
}
catch (Exception e)
{
Logger.Warn(e, "Error loading image");
notificationService.Show("Error loading image", e.Message);
}
}
///
/// Loads the user image from the given ImageSource.
///
/// The ImageSource object representing the user image.
[MethodImpl(MethodImplOptions.Synchronized)]
private void LoadUserImage(ImageSource image)
{
var current = ImageSource;
ImageSource = image;
current?.Dispose();
}
}