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

#ifndef CHROME_BROWSER_CONTEXTUAL_SEARCH_CONTEXTUAL_SEARCH_WEB_CONTENTS_HELPER_H_
#define CHROME_BROWSER_CONTEXTUAL_SEARCH_CONTEXTUAL_SEARCH_WEB_CONTENTS_HELPER_H_

#include <memory>
#include <optional>
#include <vector>

#include "base/no_destructor.h"
#include "base/uuid.h"
#include "components/contextual_search/contextual_search_session_handle.h"
#include "components/contextual_search/input_state_model.h"
#include "content/public/browser/web_contents_user_data.h"

namespace content {
class WebContents;
}

// Helper class that scopes the `ContextualSearchSessionHandle`'s
// lifetime to a `content::WebContents`. Used for transferring a contextual
// search session from one WebUI (e.g., Omnibox) to another (i.e., Co-Browsing)
// when the user submits a query.
class ContextualSearchWebContentsHelper
    : public content::WebContentsUserData<ContextualSearchWebContentsHelper> {
 public:
  ContextualSearchWebContentsHelper(const ContextualSearchWebContentsHelper&) =
      delete;
  ContextualSearchWebContentsHelper& operator=(
      const ContextualSearchWebContentsHelper&) = delete;
  ~ContextualSearchWebContentsHelper() override;

  // Sets the task ID and the contextual search session handle for the task.
  // `task_id` can be std::nullopt when transferring session before task
  // assignment.
  void SetTaskSession(
      std::optional<base::Uuid> task_id,
      std::unique_ptr<contextual_search::ContextualSearchSessionHandle> handle,
      std::unique_ptr<contextual_search::InputStateModel> input_state_model,
      std::vector<int32_t> selected_tab_ids = {}) {
    if (!selected_tab_ids.empty()) {
      selected_tab_ids_ = std::move(selected_tab_ids);
    } else if (task_id_.has_value() && task_id != task_id_) {
      // If the task ID is changing to a different non-null value and no new
      // IDs are provided, clear the old restored IDs. This happens when
      // starting a new thread in cobrowse.
      selected_tab_ids_.clear();
    }

    task_id_ = task_id;
    session_handle_ = std::move(handle);
    input_state_model_ = std::move(input_state_model);
  }
  // Returns the contextual search session handle. May return nullptr.
  contextual_search::ContextualSearchSessionHandle* session_handle() const {
    return session_handle_.get();
  }

  // Returns the input state model. May return nullptr. This transfers
  // ownership to the caller.
  std::unique_ptr<contextual_search::InputStateModel> TakeInputStateModel();

  // Returns the session handle. May return nullptr. This transfers ownership
  // to the caller.
  std::unique_ptr<contextual_search::ContextualSearchSessionHandle>
  TakeSessionHandle();

  // Returns the selected tab IDs.
  const std::vector<int32_t>& GetSelectedTabIds() const;

  // Returns the task ID associated with the current contextual search session.
  // std::nullopt if the web_contents isn't showing a contextual task.
  const std::optional<base::Uuid>& task_id() const { return task_id_; }

  // Returns contextual search session handle only if it matches `task_id`.
  // Returns nullptr if no session exists or `task_id` doesn't match.
  // This will update the task ID to be `task_id` if it was previously empty.
  contextual_search::ContextualSearchSessionHandle* GetSessionForTask(
      const base::Uuid& task_id) {
    if (!task_id_) {
      task_id_ = std::make_optional(task_id);
    }
    return (session_handle_ && task_id_ == task_id) ? session_handle_.get()
                                                    : nullptr;
  }

  std::unique_ptr<contextual_search::InputStateModel>
  TakeInputStateModelForTask(const base::Uuid& task_id) {
    if (!task_id_) {
      task_id_ = std::make_optional(task_id);
    }
    // Return and transfer ownership of the model if it matches the task.
    if (task_id_ == task_id) {
      return TakeInputStateModel();
    }
    return nullptr;
  }

  const std::vector<int32_t>& GetSelectedTabIdsForTask(
      const base::Uuid& task_id) {
    if (!task_id_) {
      task_id_ = std::make_optional(task_id);
    }

    if (task_id_ == task_id) {
      return GetSelectedTabIds();
    }
    static const base::NoDestructor<std::vector<int32_t>> empty_vector;
    return *empty_vector;
  }

 private:
  explicit ContextualSearchWebContentsHelper(
      content::WebContents* web_contents);
  friend class content::WebContentsUserData<ContextualSearchWebContentsHelper>;

  // The task ID the session handle is associated with, if any.
  std::optional<base::Uuid> task_id_;
  std::unique_ptr<contextual_search::ContextualSearchSessionHandle>
      session_handle_;
  std::unique_ptr<contextual_search::InputStateModel> input_state_model_;
  std::vector<int32_t> selected_tab_ids_;

  WEB_CONTENTS_USER_DATA_KEY_DECL();
};

#endif  // CHROME_BROWSER_CONTEXTUAL_SEARCH_CONTEXTUAL_SEARCH_WEB_CONTENTS_HELPER_H_