1  // Copyright (c) Microsoft Corporation
  2  // The Microsoft Corporation licenses this file to you under the MIT license.
  3  // See the LICENSE file in the project root for more information.
  4  
  5  using System.ComponentModel;
  6  using CommunityToolkit.Mvvm.ComponentModel;
  7  using CommunityToolkit.Mvvm.Input;
  8  using CommunityToolkit.Mvvm.Messaging;
  9  using Microsoft.CmdPal.Core.Common;
 10  using Microsoft.CmdPal.Core.ViewModels.Messages;
 11  using Microsoft.CmdPal.Core.ViewModels.Models;
 12  using Microsoft.CommandPalette.Extensions;
 13  
 14  namespace Microsoft.CmdPal.Core.ViewModels;
 15  
 16  public partial class ShellViewModel : ObservableObject,
 17      IDisposable,
 18      IRecipient<PerformCommandMessage>,
 19      IRecipient<HandleCommandResultMessage>
 20  {
 21      private readonly IRootPageService _rootPageService;
 22      private readonly IAppHostService _appHostService;
 23      private readonly TaskScheduler _scheduler;
 24      private readonly IPageViewModelFactoryService _pageViewModelFactory;
 25      private readonly Lock _invokeLock = new();
 26      private Task? _handleInvokeTask;
 27  
 28      // Cancellation token source for page loading/navigation operations
 29      private CancellationTokenSource? _navigationCts;
 30  
 31      [ObservableProperty]
 32      public partial bool IsLoaded { get; set; } = false;
 33  
 34      [ObservableProperty]
 35      public partial DetailsViewModel? Details { get; set; }
 36  
 37      [ObservableProperty]
 38      public partial bool IsDetailsVisible { get; set; }
 39  
 40      [ObservableProperty]
 41      public partial bool IsSearchBoxVisible { get; set; } = true;
 42  
 43      private PageViewModel _currentPage;
 44  
 45      public PageViewModel CurrentPage
 46      {
 47          get => _currentPage;
 48          set
 49          {
 50              var oldValue = _currentPage;
 51              if (SetProperty(ref _currentPage, value))
 52              {
 53                  oldValue.PropertyChanged -= CurrentPage_PropertyChanged;
 54                  value.PropertyChanged += CurrentPage_PropertyChanged;
 55  
 56                  if (oldValue is IDisposable disposable)
 57                  {
 58                      try
 59                      {
 60                          disposable.Dispose();
 61                      }
 62                      catch (Exception ex)
 63                      {
 64                          CoreLogger.LogError(ex.ToString());
 65                      }
 66                  }
 67              }
 68          }
 69      }
 70  
 71      private void CurrentPage_PropertyChanged(object? sender, PropertyChangedEventArgs e)
 72      {
 73          if (e.PropertyName == nameof(PageViewModel.HasSearchBox))
 74          {
 75              IsSearchBoxVisible = CurrentPage.HasSearchBox;
 76          }
 77      }
 78  
 79      private IPage? _rootPage;
 80  
 81      private bool _isNested;
 82  
 83      public bool IsNested => _isNested;
 84  
 85      public PageViewModel NullPage { get; private set; }
 86  
 87      public ShellViewModel(
 88          TaskScheduler scheduler,
 89          IRootPageService rootPageService,
 90          IPageViewModelFactoryService pageViewModelFactory,
 91          IAppHostService appHostService)
 92      {
 93          _pageViewModelFactory = pageViewModelFactory;
 94          _scheduler = scheduler;
 95          _rootPageService = rootPageService;
 96          _appHostService = appHostService;
 97  
 98          NullPage = new NullPageViewModel(_scheduler, appHostService.GetDefaultHost());
 99          _currentPage = new LoadingPageViewModel(null, _scheduler, appHostService.GetDefaultHost());
100  
101          // Register to receive messages
102          WeakReferenceMessenger.Default.Register<PerformCommandMessage>(this);
103          WeakReferenceMessenger.Default.Register<HandleCommandResultMessage>(this);
104      }
105  
106      [RelayCommand]
107      public async Task<bool> LoadAsync()
108      {
109          // First, do any loading that the root page service needs to do before we can
110          // display the root page. For example, this might include loading
111          // the built-in commands, or loading the settings.
112          await _rootPageService.PreLoadAsync();
113  
114          IsLoaded = true;
115  
116          // Now that the basics are set up, we can load the root page.
117          _rootPage = _rootPageService.GetRootPage();
118  
119          // This sends a message to us to load the root page view model.
120          WeakReferenceMessenger.Default.Send<PerformCommandMessage>(new(new ExtensionObject<ICommand>(_rootPage)));
121  
122          // Now that the root page is loaded, do any post-load work that the root page service needs to do.
123          // This runs asynchronously, on a background thread.
124          // This might include starting extensions, for example.
125          // Note: We don't await this, so that we can return immediately.
126          // This is important because we don't want to block the UI thread.
127          _ = Task.Run(async () =>
128          {
129              await _rootPageService.PostLoadRootPageAsync();
130          });
131  
132          return true;
133      }
134  
135      private async Task LoadPageViewModelAsync(PageViewModel viewModel, CancellationToken cancellationToken = default)
136      {
137          // Note: We removed the general loading state, extensions sometimes use their `IsLoading`, but it's inconsistently implemented it seems.
138          // IsInitialized is our main indicator of the general overall state of loading props/items from a page we use for the progress bar
139          // This triggers that load generally with the InitializeCommand asynchronously when we navigate to a page.
140          // We could re-track the page loading status, if we need it more granularly below again, but between the initialized and error message, we can infer some status.
141          // We could also maybe move this thread offloading we do for loading into PageViewModel and better communicate between the two... a few different options.
142  
143          ////LoadedState = ViewModelLoadedState.Loading;
144          if (!viewModel.IsInitialized
145              && viewModel.InitializeCommand is not null)
146          {
147              var outer = Task.Run(
148                  async () =>
149                  {
150                      // You know, this creates the situation where we wait for
151                      // both loading page properties, AND the items, before we
152                      // display anything.
153                      //
154                      // We almost need to do an async await on initialize, then
155                      // just a fire-and-forget on FetchItems.
156                      // RE: We do set the CurrentPage in ShellPage.xaml.cs as well, so, we kind of are doing two different things here.
157                      // Definitely some more clean-up to do, but at least its centralized to one spot now.
158                      viewModel.InitializeCommand.Execute(null);
159  
160                      await viewModel.InitializeCommand.ExecutionTask!;
161  
162                      if (viewModel.InitializeCommand.ExecutionTask.Status != TaskStatus.RanToCompletion)
163                      {
164                          if (viewModel.InitializeCommand.ExecutionTask.Exception is AggregateException ex)
165                          {
166                              CoreLogger.LogError(ex.ToString());
167                          }
168                      }
169                      else
170                      {
171                          var t = Task.Factory.StartNew(
172                              () =>
173                              {
174                                  if (cancellationToken.IsCancellationRequested)
175                                  {
176                                      if (viewModel is IDisposable disposable)
177                                      {
178                                          try
179                                          {
180                                              disposable.Dispose();
181                                          }
182                                          catch (Exception ex)
183                                          {
184                                              CoreLogger.LogError(ex.ToString());
185                                          }
186                                      }
187  
188                                      return;
189                                  }
190  
191                                  CurrentPage = viewModel;
192                              },
193                              cancellationToken,
194                              TaskCreationOptions.None,
195                              _scheduler);
196                          await t;
197                      }
198                  },
199                  cancellationToken);
200              await outer;
201          }
202          else
203          {
204              if (cancellationToken.IsCancellationRequested)
205              {
206                  if (viewModel is IDisposable disposable)
207                  {
208                      try
209                      {
210                          disposable.Dispose();
211                      }
212                      catch (Exception ex)
213                      {
214                          CoreLogger.LogError(ex.ToString());
215                      }
216                  }
217  
218                  return;
219              }
220  
221              CurrentPage = viewModel;
222          }
223      }
224  
225      public void Receive(PerformCommandMessage message)
226      {
227          PerformCommand(message);
228      }
229  
230      private void PerformCommand(PerformCommandMessage message)
231      {
232          // Create/replace the navigation cancellation token.
233          // If one already exists, cancel and dispose it first.
234          var newCts = new CancellationTokenSource();
235          var oldCts = Interlocked.Exchange(ref _navigationCts, newCts);
236          if (oldCts is not null)
237          {
238              try
239              {
240                  oldCts.Cancel();
241              }
242              catch (Exception ex)
243              {
244                  CoreLogger.LogError(ex.ToString());
245              }
246              finally
247              {
248                  oldCts.Dispose();
249              }
250          }
251  
252          var navigationToken = newCts.Token;
253  
254          var command = message.Command.Unsafe;
255          if (command is null)
256          {
257              return;
258          }
259  
260          var host = _appHostService.GetHostForCommand(message.Context, CurrentPage.ExtensionHost);
261  
262          _rootPageService.OnPerformCommand(message.Context, !CurrentPage.IsNested, host);
263  
264          try
265          {
266              if (command is IPage page)
267              {
268                  CoreLogger.LogDebug($"Navigating to page");
269  
270                  var isMainPage = command == _rootPage;
271                  _isNested = !isMainPage;
272  
273                  // Telemetry: Track extension page navigation for session metrics
274                  if (host is not null)
275                  {
276                      string extensionId = host.GetExtensionDisplayName() ?? "builtin";
277                      string commandId = command?.Id ?? "unknown";
278                      string commandName = command?.Name ?? "unknown";
279                      WeakReferenceMessenger.Default.Send<TelemetryExtensionInvokedMessage>(
280                          new(extensionId, commandId, commandName, true, 0));
281                  }
282  
283                  // Construct our ViewModel of the appropriate type and pass it the UI Thread context.
284                  var pageViewModel = _pageViewModelFactory.TryCreatePageViewModel(page, _isNested, host!);
285                  if (pageViewModel is null)
286                  {
287                      CoreLogger.LogError($"Failed to create ViewModel for page {page.GetType().Name}");
288                      throw new NotSupportedException();
289                  }
290  
291                  // Clear command bar, ViewModel initialization can already set new commands if it wants to
292                  OnUIThread(() => WeakReferenceMessenger.Default.Send<UpdateCommandBarMessage>(new(null)));
293  
294                  // Kick off async loading of our ViewModel
295                  LoadPageViewModelAsync(pageViewModel, navigationToken)
296                      .ContinueWith(
297                          (Task t) =>
298                          {
299                              // clean up the navigation token if it's still ours
300                              if (Interlocked.CompareExchange(ref _navigationCts, null, newCts) == newCts)
301                              {
302                                  newCts.Dispose();
303                              }
304                          },
305                          navigationToken,
306                          TaskContinuationOptions.None,
307                          _scheduler);
308  
309                  // While we're loading in the background, immediately move to the next page.
310                  WeakReferenceMessenger.Default.Send<NavigateToPageMessage>(new(pageViewModel, message.WithAnimation, navigationToken));
311  
312                  // Note: Originally we set our page back in the ViewModel here, but that now happens in response to the Frame navigating triggered from the above
313                  // See RootFrame_Navigated event handler.
314              }
315              else if (command is IInvokableCommand invokable)
316              {
317                  CoreLogger.LogDebug($"Invoking command");
318  
319                  WeakReferenceMessenger.Default.Send<TelemetryBeginInvokeMessage>();
320                  StartInvoke(message, invokable, host);
321              }
322          }
323          catch (Exception ex)
324          {
325              // TODO: It would be better to do this as a page exception, rather
326              // than a silent log message.
327              host?.Log(ex.Message);
328          }
329      }
330  
331      private void StartInvoke(PerformCommandMessage message, IInvokableCommand invokable, AppExtensionHost? host)
332      {
333          // TODO GH #525 This needs more better locking.
334          lock (_invokeLock)
335          {
336              if (_handleInvokeTask is not null)
337              {
338                  // do nothing - a command is already doing a thing
339              }
340              else
341              {
342                  _handleInvokeTask = Task.Run(() =>
343                  {
344                      SafeHandleInvokeCommandSynchronous(message, invokable, host);
345                  });
346              }
347          }
348      }
349  
350      private void SafeHandleInvokeCommandSynchronous(PerformCommandMessage message, IInvokableCommand invokable, AppExtensionHost? host)
351      {
352          // Telemetry: Track command execution time and success
353          var stopwatch = System.Diagnostics.Stopwatch.StartNew();
354          var command = message.Command.Unsafe;
355          string extensionId = host?.GetExtensionDisplayName() ?? "builtin";
356          string commandId = command?.Id ?? "unknown";
357          string commandName = command?.Name ?? "unknown";
358          bool success = false;
359  
360          try
361          {
362              // Call out to extension process.
363              // * May fail!
364              // * May never return!
365              var result = invokable.Invoke(message.Context);
366  
367              // But if it did succeed, we need to handle the result.
368              UnsafeHandleCommandResult(result);
369  
370              success = true;
371              _handleInvokeTask = null;
372          }
373          catch (Exception ex)
374          {
375              success = false;
376              _handleInvokeTask = null;
377  
378              // Telemetry: Track errors for session metrics
379              WeakReferenceMessenger.Default.Send<ErrorOccurredMessage>(new());
380  
381              // TODO: It would be better to do this as a page exception, rather
382              // than a silent log message.
383              host?.Log(ex.Message);
384          }
385          finally
386          {
387              // Telemetry: Send extension invocation metrics (always sent, even on failure)
388              stopwatch.Stop();
389              WeakReferenceMessenger.Default.Send<TelemetryExtensionInvokedMessage>(
390                  new(extensionId, commandId, commandName, success, (ulong)stopwatch.ElapsedMilliseconds));
391          }
392      }
393  
394      private void UnsafeHandleCommandResult(ICommandResult? result)
395      {
396          if (result is null)
397          {
398              // No result, nothing to do.
399              return;
400          }
401  
402          var kind = result.Kind;
403          CoreLogger.LogDebug($"handling {kind.ToString()}");
404  
405          WeakReferenceMessenger.Default.Send<TelemetryInvokeResultMessage>(new(kind));
406          switch (kind)
407          {
408              case CommandResultKind.Dismiss:
409                  {
410                      // Reset the palette to the main page and dismiss
411                      GoHome(withAnimation: false, focusSearch: false);
412                      WeakReferenceMessenger.Default.Send(new DismissMessage());
413                      break;
414                  }
415  
416              case CommandResultKind.GoHome:
417                  {
418                      // Go back to the main page, but keep it open
419                      GoHome();
420                      break;
421                  }
422  
423              case CommandResultKind.GoBack:
424                  {
425                      GoBack();
426                      break;
427                  }
428  
429              case CommandResultKind.Hide:
430                  {
431                      // Keep this page open, but hide the palette.
432                      WeakReferenceMessenger.Default.Send(new DismissMessage());
433                      break;
434                  }
435  
436              case CommandResultKind.KeepOpen:
437                  {
438                      // Do nothing.
439                      break;
440                  }
441  
442              case CommandResultKind.Confirm:
443                  {
444                      if (result.Args is IConfirmationArgs a)
445                      {
446                          WeakReferenceMessenger.Default.Send<ShowConfirmationMessage>(new(a));
447                      }
448  
449                      break;
450                  }
451  
452              case CommandResultKind.ShowToast:
453                  {
454                      if (result.Args is IToastArgs a)
455                      {
456                          WeakReferenceMessenger.Default.Send<ShowToastMessage>(new(a.Message));
457                          UnsafeHandleCommandResult(a.Result);
458                      }
459  
460                      break;
461                  }
462          }
463      }
464  
465      public void GoHome(bool withAnimation = true, bool focusSearch = true)
466      {
467          _rootPageService.GoHome();
468          WeakReferenceMessenger.Default.Send<GoHomeMessage>(new(withAnimation, focusSearch));
469      }
470  
471      public void GoBack(bool withAnimation = true, bool focusSearch = true)
472      {
473          WeakReferenceMessenger.Default.Send<GoBackMessage>(new(withAnimation, focusSearch));
474      }
475  
476      public void Receive(HandleCommandResultMessage message)
477      {
478          UnsafeHandleCommandResult(message.Result.Unsafe);
479      }
480  
481      private void OnUIThread(Action action)
482      {
483          _ = Task.Factory.StartNew(
484              action,
485              CancellationToken.None,
486              TaskCreationOptions.None,
487              _scheduler);
488      }
489  
490      public void CancelNavigation()
491      {
492          _navigationCts?.Cancel();
493      }
494  
495      public void Dispose()
496      {
497          _handleInvokeTask?.Dispose();
498          _navigationCts?.Dispose();
499  
500          GC.SuppressFinalize(this);
501      }
502  }