// 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 COMPONENTS_PAGE_CONTENT_ANNOTATIONS_CONTENT_PAGE_CONTEXT_FETCHER_H_
#define COMPONENTS_PAGE_CONTENT_ANNOTATIONS_CONTENT_PAGE_CONTEXT_FETCHER_H_

#include <cstdint>
#include <memory>
#include <optional>
#include <string>
#include <variant>
#include <vector>

#include "base/feature_list.h"
#include "base/functional/callback_forward.h"
#include "base/functional/callback_helpers.h"
#include "base/gtest_prod_util.h"
#include "base/memory/weak_ptr.h"
#include "base/metrics/field_trial_params.h"
#include "base/task/single_thread_task_runner.h"
#include "base/time/time.h"
#include "base/timer/elapsed_timer.h"
#include "base/types/expected.h"
#include "components/content_extraction/content/browser/inner_text.h"
#include "components/optimization_guide/content/browser/page_content_proto_provider.h"
#include "components/optimization_guide/proto/features/common_quality_data.pb.h"
#include "components/page_content_annotations/content/page_context_fetcher_options.h"
#include "components/viz/common/surfaces/tracked_element_rects.h"
#include "content/public/browser/render_widget_host_view.h"
#include "content/public/browser/web_contents_observer.h"
#include "pdf/buildflags.h"
#include "third_party/skia/include/core/SkBitmap.h"
#include "third_party/skia/include/core/SkColor.h"
#include "ui/gfx/geometry/rect.h"
#include "ui/gfx/geometry/size.h"
#include "url/origin.h"

#if BUILDFLAG(ENABLE_PDF)
#include "pdf/mojom/pdf.mojom.h"
#endif  // BUILDFLAG(ENABLE_PDF)

namespace content {
class BrowserContext;
class WebContents;
}  // namespace content

namespace page_content_annotations {

class PageContentScreenshotService;

// TODO(b/504577535): Support PDF bookmark extraction.
// TODO(b/504577256): Support PDF accessibility info extraction.
struct PdfResult {
  explicit PdfResult(url::Origin origin);
  PdfResult(url::Origin origin, std::vector<uint8_t> bytes);
  PdfResult(url::Origin origin, std::string text);
  PdfResult(const PdfResult&) = delete;
  PdfResult& operator=(const PdfResult&) = delete;
  PdfResult(PdfResult&&);
  PdfResult& operator=(PdfResult&&);
  ~PdfResult();

  url::Origin origin;

  // The PDF extraction result can be either bytes or string, depending on which
  // extraction option is selected.
  std::variant<std::vector<uint8_t>, std::string> data;

  bool size_exceeded = false;
};

struct ScreenshotResult {
  explicit ScreenshotResult(gfx::Size dimensions);
  ScreenshotResult(const ScreenshotResult&) = delete;
  ScreenshotResult& operator=(const ScreenshotResult&) = delete;
  ScreenshotResult(ScreenshotResult&&);
  ScreenshotResult& operator=(ScreenshotResult&&);
  ~ScreenshotResult();
  std::vector<uint8_t> screenshot_data;
  std::string mime_type;
  gfx::Size dimensions;
  base::TimeTicks end_time;
};

struct InnerTextResultWithTruncation
    : public content_extraction::InnerTextResult {
  InnerTextResultWithTruncation(std::string inner_text,
                                std::optional<unsigned> node_offset,
                                bool truncated);
  ~InnerTextResultWithTruncation();
  bool truncated = false;
};

struct PageContentResultWithEndTime
    : public optimization_guide::AIPageContentResult {
  explicit PageContentResultWithEndTime(
      optimization_guide::AIPageContentResult&& result);
  base::TimeTicks end_time;
};

struct FetchPageContextResult {
  FetchPageContextResult();
  FetchPageContextResult(const FetchPageContextResult&) = delete;
  FetchPageContextResult& operator=(const FetchPageContextResult&) = delete;
  FetchPageContextResult(FetchPageContextResult&&);
  FetchPageContextResult& operator=(FetchPageContextResult&&);
  ~FetchPageContextResult();
  base::expected<ScreenshotResult, std::string> screenshot_result;
  std::optional<InnerTextResultWithTruncation> inner_text_result;
  std::optional<PdfResult> pdf_result;
  base::expected<PageContentResultWithEndTime, std::string>
      annotated_page_content_result;
  std::optional<optimization_guide::proto::ScreenshotInfo> screenshot_info;
};

enum class FetchPageContextError {
  kUnknown,
  kWebContentsChanged,
  // The context is not eligible for sharing.
  kPageContextNotEligible,
  kWebContentsWentAway,
};

std::string ToString(FetchPageContextError error);

// TODO(bokan): message is redundant with error_code. Replace usage with
// ToString.
struct FetchPageContextErrorDetails {
  FetchPageContextError error_code = FetchPageContextError::kUnknown;
  std::string message;
};
using FetchPageContextResultCallbackArg =
    base::expected<std::unique_ptr<FetchPageContextResult>,
                   FetchPageContextErrorDetails>;

// Controls scaling and quality of tab screenshots.
// Does not override screenshot_collection_options if they are set, only
// modifies the default values.
BASE_DECLARE_FEATURE(kGlicTabScreenshotExperiment);

// Controls whether password fields are redacted from screenshots.
BASE_DECLARE_FEATURE(kGlicScreenshotPasswordRedaction);

// Controls whether sensitive payment fields are redacted from screenshots.
BASE_DECLARE_FEATURE(kGlicScreenshotSensitivePaymentRedaction);

extern const base::FeatureParam<int> kMaxScreenshotWidthParam;

extern const base::FeatureParam<int> kMaxScreenshotHeightParam;

extern const base::FeatureParam<int> kScreenshotQuality;

extern const base::FeatureParam<std::string> kScreenshotImageType;

extern const base::FeatureParam<base::TimeDelta> kScreenshotTimeout;

extern const base::FeatureParam<base::TimeDelta>
    kScreenshotTimeoutBrowserAllowance;

// Callback used for relaying progress.
class FetchPageProgressListener {
 public:
  virtual ~FetchPageProgressListener() = default;
  virtual void BeginScreenshot() = 0;
  virtual void ScreenshotCaptured(const SkBitmap& bitmap) = 0;
  virtual void ScreenshotRedacted(const SkBitmap& bitmap) = 0;
  virtual void EndScreenshot(std::optional<std::string> error) = 0;
  virtual void BeginAPC() = 0;
  virtual void EndAPC(std::optional<std::string> error) = 0;
};

using FetchPageContextResultCallback =
    base::OnceCallback<void(FetchPageContextResultCallbackArg)>;

using GetScreenshotServiceCallback =
    base::RepeatingCallback<PageContentScreenshotService*(
        content::BrowserContext*)>;

// Encodes a screenshot according to the enabled feature flags.
std::optional<std::vector<uint8_t>> EncodeScreenshot(const SkBitmap& bitmap,
    const std::optional<ScreenshotOptions::ScreenshotCollectionOptions>&
        screenshot_collection_options);

// Coordinates fetching multiple types of page context.
class PageContextFetcher : public content::WebContentsObserver {
 public:
  explicit PageContextFetcher(
      GetScreenshotServiceCallback get_screenshot_service_callback,
      std::unique_ptr<FetchPageProgressListener> progress_listener);
  ~PageContextFetcher() override;

  void FetchStart(content::WebContents& aweb_contents,
                  const FetchPageContextOptions& options,
                  FetchPageContextResultCallback callback);

 private:
  FRIEND_TEST_ALL_PREFIXES(PageContextFetcherTest,
                           RedactScreenshotOnWorkerThread);
  FRIEND_TEST_ALL_PREFIXES(PageContextFetcherTest,
                           RedactScreenshotOnWorkerThreadNoRedaction);
  FRIEND_TEST_ALL_PREFIXES(PageContextFetcherIframeInfoTest,
                           AddIframeInfoSuccess);
  FRIEND_TEST_ALL_PREFIXES(PageContextFetcherIframeInfoTest,
                           AddIframeInfoNoUrlOrigin);
  FRIEND_TEST_ALL_PREFIXES(PageContextFetcherIframeInfoTest,
                           NoIframeInfoWhenFeatureDisabled);

  // Redacts a screenshot by painting over sensitive regions with
  // `redaction_color`.
  static base::expected<SkBitmap, std::string> RedactScreenshotOnWorkerThread(
      const SkBitmap& bitmap,
      const std::vector<gfx::Rect>& visible_bounding_boxes_for_redaction,
      SkColor4f redaction_color);

#if BUILDFLAG(ENABLE_PDF)
  void FetchPdfContent(const PdfOptions& options);
  void ReceivedPdfBytes(url::Origin pdf_origin,
                        uint32_t pdf_size_limit,
                        pdf::mojom::PdfListener::GetPdfBytesStatus status,
                        const std::vector<uint8_t>& pdf_bytes,
                        uint32_t page_count);
  void ReceivedPdfText(url::Origin pdf_origin,
                       uint32_t text_byte_limit,
                       const std::u16string& text);
#endif  // BUILDFLAG(ENABLE_PDF)

  void GetTabScreenshot(content::WebContents& web_contents,
                        const ScreenshotOptions& screenshot_options);

  void SetCaptureCountLock(content::WebContents& web_contents);

  void ScheduleScreenshotTimeout();

  void ReceivedViewportBitmap(const content::CopyFromSurfaceResult& result);

  void RedactAndEncodeScreenshot(
      std::vector<gfx::Rect> visible_bounding_boxes_for_redaction);

  void RedactAndEncodeScreenshotIfNeeded();

  void ReceivedViewportBitmapOrError(
      const viz::TrackedElementRects& tracked_element_rects,
      base::expected<const SkBitmap*, std::string> bitmap_result);

  // content::WebContentsObserver impl.
  void PrimaryPageChanged(content::Page& page) override;

  void OnScreenshotTimeout();

  void ReceivedEncodedScreenshot(
      base::expected<std::pair<std::vector<uint8_t>, SkBitmap>, std::string>
          screenshot_data);

  void ReceivedInnerText(
      std::unique_ptr<content_extraction::InnerTextResult> result);

  void ReceivedAnnotatedPageContent(
      optimization_guide::AIPageContentResultOrError content);

  void RunCallbackIfComplete();

  void ProcessTrackedElementRects(
      const viz::TrackedElementRects& tracked_element_rects);

  void CollectTrackedElementRectsForIframes(
      const viz::TrackedElementRects& tracked_element_rects);

  void MaybeAddIframeInfo();

  void CollectTrackedElementRectsForPassword(
      const viz::TrackedElementRects& tracked_element_rects);

  base::WeakPtr<PageContextFetcher> GetWeakPtr();

  const GetScreenshotServiceCallback get_screenshot_service_callback_;
  FetchPageContextResultCallback callback_;

  uint32_t inner_text_bytes_limit_ = 0;

  // screenshot processing dependencies.
  std::optional<SkBitmap> screenshot_bitmap_;
  bool screenshot_needs_redaction_using_apc_ = false;
  std::vector<gfx::Rect> tracked_element_bounds_for_screenshot_redaction_;
  gfx::Size original_view_size_pixels_;

  std::vector<optimization_guide::proto::IframeInfo> iframe_info_;

  // Intermediate results:

  // Whether work is complete for each task, does not imply success.
  bool initialization_done_ = false;
  bool screenshot_capture_done_ = false;
  bool screenshot_done_ = false;
  bool inner_text_done_ = false;
  bool pdf_done_ = false;
  bool annotated_page_content_done_ = false;
  SkColor4f screenshot_redaction_color_ = SkColors::kBlack;
  std::optional<ScreenshotOptions::ScreenshotCollectionOptions>
      screenshot_collection_options_;
  // Whether the primary page has changed since context fetching began.
  bool primary_page_changed_ = false;
  std::unique_ptr<FetchPageContextResult> pending_result_;
  base::ElapsedTimer elapsed_timer_;
  base::ScopedClosureRunner capture_count_lock_;

  std::unique_ptr<FetchPageProgressListener> progress_listener_;

  base::WeakPtrFactory<PageContextFetcher> weak_ptr_factory_{this};
};

}  // namespace page_content_annotations

#endif  // COMPONENTS_PAGE_CONTENT_ANNOTATIONS_CONTENT_PAGE_CONTEXT_FETCHER_H_