// Copyright 2022 The Chromium Authors
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

#include "ash/wm/desks/templates/restore_data_collector.h"

#include "ash/multi_user/multi_user_window_manager.h"
#include "ash/public/cpp/saved_desk_delegate.h"
#include "ash/public/cpp/window_properties.h"
#include "ash/shell.h"
#include "ash/wm/desks/desk.h"
#include "ash/wm/desks/templates/saved_desk_dialog_controller.h"
#include "ash/wm/desks/templates/saved_desk_util.h"
#include "ash/wm/mru_window_tracker.h"
#include "ash/wm/overview/overview_controller.h"
#include "ash/wm/window_restore/window_restore_util.h"
#include "base/barrier_callback.h"
#include "base/uuid.h"
#include "chromeos/constants/chromeos_features.h"
#include "components/app_restore/app_launch_info.h"
#include "components/app_restore/restore_data.h"
#include "components/app_restore/window_info.h"
#include "ui/wm/core/window_util.h"

namespace ash {

RestoreDataCollector::Call::Call()
    : data(std::make_unique<app_restore::RestoreData>()) {}
RestoreDataCollector::Call::~Call() = default;

RestoreDataCollector::WindowData::WindowData() = default;
RestoreDataCollector::WindowData::WindowData(WindowData&&) = default;
RestoreDataCollector::WindowData& RestoreDataCollector::WindowData::operator=(
    WindowData&&) = default;
RestoreDataCollector::WindowData::~WindowData() = default;

RestoreDataCollector::RestoreDataCollector() = default;
RestoreDataCollector::~RestoreDataCollector() = default;

void RestoreDataCollector::CaptureActiveDeskAsSavedDesk(
    GetDeskTemplateCallback callback,
    DeskTemplateType template_type,
    const std::string& template_name,
    aura::Window* root_window_to_show,
    AccountId current_account_id,
    const base::flat_set<std::string>& coral_app_id_allowlist) {
  DCHECK_CALLED_ON_VALID_SEQUENCE(sequence_checker_);

  auto call = std::make_unique<Call>();
  call->root_window_to_show = root_window_to_show;
  call->template_type = template_type;
  call->template_name = template_name;
  auto* window_manager = MultiUserWindowManager::Get();
  auto* const shell = Shell::Get();
  auto mru_windows =
      shell->mru_window_tracker()->BuildMruWindowList(kActiveDesk);
  auto* delegate = shell->saved_desk_delegate();

  std::vector<aura::Window*> supported_windows;

  for (aura::Window* window : mru_windows) {
    // Skip transient windows without reporting.
    if (wm::GetTransientParent(window)) {
      continue;
    }

    // Coral desk templates only contain a subset of the apps on the active
    // desk.
    if (template_type == DeskTemplateType::kCoral &&
        !coral_app_id_allowlist.contains(*window->GetProperty(kAppIDKey))) {
      continue;
    }

    // If `window_manager` is not nullptr, then we have a multi profile
    // session. We need to make sure that the windows we are capturing belongs
    // to the owner of the current active session.
    if (window_manager) {
      // Skip windows that belong to another profile user.
      const AccountId& window_owner = window_manager->GetWindowOwner(window);
      if (window_owner.is_valid() && current_account_id != window_owner) {
        continue;
      }
    }

    if (!delegate->IsWindowSupportedForSavedDesk(window)) {
      call->unsupported_apps.push_back(window);
      if (!delegate->IsWindowPersistable(window)) {
        ++call->non_persistable_window_count;
      }
      continue;
    }

    // Skip windows that do not associate with a full restore app id.
    const std::string app_id = saved_desk_util::GetAppId(window);
    if (!Shell::Get()
             ->overview_controller()
             ->disable_app_id_check_for_saved_desks() &&
        app_id.empty()) {
      call->unsupported_apps.push_back(window);
      continue;
    }

    supported_windows.push_back(window);
  }

  // Do not create a saved desk if the desk is empty or only contains
  // unsupported apps, unless it's a Floating Workspace desk.
  if (supported_windows.empty() &&
      template_type != DeskTemplateType::kFloatingWorkspace) {
    std::move(callback).Run(nullptr);
    return;
  }

  if (root_window_to_show) {
    window_tracker_.Add(root_window_to_show);
  }
  call->callback = std::move(callback);

  if (supported_windows.empty()) {
    SendDeskTemplate(std::move(call));
    return;
  }

  auto barrier = base::BarrierCallback<WindowData>(
      supported_windows.size(),
      base::BindOnce(&RestoreDataCollector::OnAllAppLaunchDataReceived,
                     weak_factory_.GetWeakPtr(), std::move(call)));

  for (aura::Window* window : supported_windows) {
    const std::string app_id = saved_desk_util::GetAppId(window);
    std::unique_ptr<app_restore::WindowInfo> window_info =
        BuildWindowInfo(window, /*activation_index=*/std::nullopt, mru_windows);

    // Clear the desk ID and uuid in the WindowInfo that is to be stored in
    // the template. They will be set to the newly created desk when
    // launching.
    window_info->desk_id.reset();
    window_info->desk_guid = base::Uuid();

    delegate->GetAppLaunchDataForSavedDesk(
        window, base::BindOnce(&RestoreDataCollector::OnAppLaunchDataReceived,
                               weak_factory_.GetWeakPtr(), app_id,
                               std::move(window_info), barrier));
  }
}

void RestoreDataCollector::OnAppLaunchDataReceived(
    std::string app_id,
    std::unique_ptr<app_restore::WindowInfo> window_info,
    base::RepeatingCallback<void(WindowData)> barrier,
    std::unique_ptr<app_restore::AppLaunchInfo> app_launch_info) {
  DCHECK_CALLED_ON_VALID_SEQUENCE(sequence_checker_);

  WindowData data;
  data.app_id = std::move(app_id);
  data.window_info = std::move(window_info);
  data.app_launch_info = std::move(app_launch_info);
  barrier.Run(std::move(data));
}

void RestoreDataCollector::OnAllAppLaunchDataReceived(
    std::unique_ptr<Call> call,
    std::vector<WindowData> window_data_list) {
  DCHECK_CALLED_ON_VALID_SEQUENCE(sequence_checker_);

  for (auto& window_data : window_data_list) {
    // nullptr means that this app does not have any data to save.
    if (window_data.app_launch_info) {
      const int32_t window_id = *window_data.app_launch_info->window_id;
      call->data->AddAppLaunchInfo(std::move(window_data.app_launch_info));
      call->data->ModifyWindowInfo(window_data.app_id, window_id,
                                   *window_data.window_info);
    }
  }

  SendDeskTemplate(std::move(call));
}

void RestoreDataCollector::SendDeskTemplate(std::unique_ptr<Call> call) {
  DCHECK_CALLED_ON_VALID_SEQUENCE(sequence_checker_);

  auto desk_template = std::make_unique<DeskTemplate>(
      base::Uuid::GenerateRandomV4(), DeskTemplateSource::kUser,
      call->template_name, base::Time::Now(), call->template_type);
  desk_template->set_desk_restore_data(std::move(call->data));

  if (!call->unsupported_apps.empty() &&
      Shell::Get()->overview_controller()->InOverviewSession()) {
    // The ideal root window may have gone by now.  In that case fall back to
    // the primary root one.
    auto* root_window_to_show = call->root_window_to_show.get();
    if (root_window_to_show && window_tracker_.Contains(root_window_to_show)) {
      window_tracker_.Remove(root_window_to_show);
    } else {
      root_window_to_show = Shell::Get()->GetPrimaryRootWindow();
    }

    // There were some unsupported apps in the active desk so open up a dialog
    // to let the user know. The dialog controller should always be available
    // here since we have already determined that we are in overview mode.
    auto* dialog_controller = saved_desk_util::GetSavedDeskDialogController();
    DCHECK(dialog_controller);
    dialog_controller->ShowUnsupportedAppsDialog(
        root_window_to_show, call->unsupported_apps,
        call->non_persistable_window_count, std::move(call->callback),
        std::move(desk_template));
  } else {
    std::move(call->callback).Run(std::move(desk_template));
  }
}

}  // namespace ash