livia-test/Livia/Models/DataBlockLoader.cs

using Livia.Models.Data;
using Livia.ViewModels;
using Livia.Views.Utility;
using System.Collections.Concurrent;
using System.IO;
using System.Text.Json;
using System.Windows;
using CommunityToolkit.Mvvm.ComponentModel;
using Microsoft.Extensions.Logging;
using CommunityToolkit.Mvvm.Messaging;
using Livia.Utility;

namespace Livia.Models;

public interface IDataBlockLoader
{
    DataBlock? CurrentDisplayingDataBlock { get; }
    bool PlaceholderVisible { get; }
    T GetSharedData<T>(string key);
    void AddViewModel(ILiviaModuleViewModel module);
    void AddDeferredLoadingModuleViewModel(ILiviaModuleViewModel module);
    Task LoadData(DataBlock dataBlock);
    Task LoadServerData(ServerJobData data);
    void Init();
    void SaveSharedData(string key, object data);
    bool SharedDataContainsKey(string key);
}

internal class DataBlockLoader : ObservableObject, IDataBlockLoader, IRecipient<LogoutMessage>
{
    public bool PlaceholderVisible { get => _placeholderVisible; protected internal set => SetProperty(ref _placeholderVisible, value); }
    public DataBlock? CurrentDisplayingDataBlock { get; private set; }

    private readonly IWarningSystem _warningSystem;
    private readonly ILogger _logger;
    private readonly IPatientInfoManager _patientInfoManager;
    private readonly IServerHandler _serverHandler;
    private readonly IRoiLegendManager _roiLegendManager;

    private readonly ConcurrentDictionary<string, object> _sharedData = new();
    private readonly List<ILiviaModuleViewModel> _moduleViewModels = [];
    private readonly List<ILiviaModuleViewModel> _deferredLoadingModuleViewModels = [];
    private readonly Dictionary<string, DataBlock> _dataBlocks = new();
    private bool _placeholderVisible = true;

    public DataBlockLoader(IWarningSystem warningSystem, IPatientInfoManager patientInfoManager, ILogger logger, IServerHandler serverHandler, IRoiLegendManager roiLegendManager)
    {
        _warningSystem = warningSystem;
        _patientInfoManager = patientInfoManager;
        _logger = logger;
        _serverHandler = serverHandler;
        _roiLegendManager = roiLegendManager;
        WeakReferenceMessenger.Default.RegisterAll(this);
    }

    public void SaveSharedData(string key, object data)
    {
        _sharedData[key] = data;
    }

    public T GetSharedData<T>(string key)
    {
        return _sharedData[key] is T t ? t : throw new KeyNotFoundException();
    }

    public bool SharedDataContainsKey(string key)
    {
        return _sharedData.ContainsKey(key);
    }

    public void AddViewModel(ILiviaModuleViewModel module)
    {
        _moduleViewModels.Add(module);
    }

    public void AddDeferredLoadingModuleViewModel(ILiviaModuleViewModel module)
    {
        _deferredLoadingModuleViewModels.Add(module);
    }

    public async Task LoadData(DataBlock dataBlock)
    {
        _serverHandler.Processing = true;
        DateTime startTime = DateTime.Now;
        try
        {
            if (string.IsNullOrEmpty(dataBlock.Key))
                throw new NullReferenceException();

            _dataBlocks[dataBlock.Key] = dataBlock;
            CurrentDisplayingDataBlock = dataBlock;
            string path = dataBlock.ResultPath;

            _sharedData.Clear();
            _roiLegendManager.Clear();

            AdditionalInfoJson additionalInfo = JsonSerializer.Deserialize<AdditionalInfoJson>(await File.ReadAllTextAsync(Path.Combine(path, "additional_info.json")));
            _patientInfoManager.Init(additionalInfo);
            SaveSharedData("additionalInfo", additionalInfo);

            string defaultMaskToLoad = additionalInfo.AtlasInfoListDict.Keys.First(k => k.Contains("cbf", StringComparison.CurrentCultureIgnoreCase));
            ImageRotationViewerControlViewModel.LoadAtlasName = defaultMaskToLoad;
            PerfusionDataGridControlViewModel.LoadAtlasName = defaultMaskToLoad;
            
            Dictionary<string, ConcurrentBag<RoiExpanderControlViewModel>> roiCollectionDictionary = new();
            SaveSharedData("roiCollection", roiCollectionDictionary);

            await Task.WhenAll(_moduleViewModels.Select(module => module.LoadData(dataBlock)));
            await Task.WhenAll(_deferredLoadingModuleViewModels.Select(module => module.LoadData(dataBlock)));
            CurrentDisplayingDataBlock = dataBlock;
            WeakReferenceMessenger.Default.Send(new DataLoadedMessage());

            PlaceholderVisible = false;
            _warningSystem.ShowDialog(WarningWindowKind.Info, false, "DataImportComplete");
        }
        catch (Exception e)
        {
            PlaceholderVisible = true;
            CurrentDisplayingDataBlock = null;
            _logger.LogError(e, "Error loading data from {path}", dataBlock.Key);
            _warningSystem.ShowDialog(WarningWindowKind.Warning, true, "InvalidDataError");
        }
        finally
        {
            _serverHandler.Processing = false;
            DateTime endTime = DateTime.Now;
            TimeSpan timeUsed = endTime - startTime;
            _logger.LogInformation("Load time = {time}", timeUsed);
        }
    }

    public async Task LoadServerData(ServerJobData data)
    {
        if (string.IsNullOrEmpty(data.Key))
        {
            _logger.LogError("Cannot load data with empty key");
            return;
        }

        if (!_dataBlocks.TryGetValue(data.Key, out DataBlock? value))
        {
            value = new DataBlock(string.Empty)
            {
                Key = data.Key
            };
            _dataBlocks[data.Key] = value;
        }
        DataBlock dataBlock = value;

        (bool success, string messageIndex) = await _serverHandler.DownloadData(dataBlock);
        if (!success)
        {
            _warningSystem.ShowDialog(WarningWindowKind.Warning, true, messageIndex);
            return;
        }

        await Application.Current.Dispatcher.Invoke(() => LoadData(dataBlock));
    }

    public void Init()
    {
        foreach (ILiviaModuleViewModel liviaModuleViewModel in _moduleViewModels)
        {
            liviaModuleViewModel.Init();
        }

        foreach (ILiviaModuleViewModel liviaModuleViewModel in _deferredLoadingModuleViewModels)
        {
            liviaModuleViewModel.Init();
        }
    }

    public void Receive(LogoutMessage message)
    {
        PlaceholderVisible = true;
    }
}