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

#include <cstdint>

#include "base/containers/enum_set.h"
#include "base/functional/callback.h"
#include "base/memory/stack_allocated.h"
#include "base/types/pass_key.h"
#include "mojo/public/cpp/bindings/lib/validation_context.h"
#include "mojo/public/cpp/bindings/pending_receiver.h"
#include "third_party/blink/public/mojom/content_extraction/ai_page_content.mojom-blink.h"
#include "third_party/blink/renderer/core/dom/document.h"
#include "third_party/blink/renderer/core/frame/local_frame_view.h"
#include "third_party/blink/renderer/modules/content_extraction/paid_content.h"
#include "third_party/blink/renderer/modules/modules_export.h"
#include "third_party/blink/renderer/platform/heap/collection_support/heap_hash_map.h"
#include "third_party/blink/renderer/platform/mojo/heap_mojo_receiver_set.h"
#include "third_party/blink/renderer/platform/mojo/heap_mojo_wrapper_mode.h"
#include "third_party/blink/renderer/platform/supplementable.h"
#include "third_party/blink/renderer/platform/wtf/hash_map.h"
#include "third_party/blink/renderer/platform/wtf/hash_set.h"
#include "third_party/blink/renderer/platform/wtf/hash_traits.h"

namespace blink {
class Document;
class LayoutIFrame;
class LayoutBox;
class LayoutObject;
class LocalFrame;
class Node;
#if DCHECK_IS_ON()
class AutoBuildHelper;
#endif

// AIPageContent is responsible for handling requests for inner-text. It calls
// to InnerTextBuilder to handle building of the text.
class MODULES_EXPORT AIPageContentAgent final
    : public GarbageCollected<AIPageContentAgent>,
      public mojom::blink::AIPageContentAgent,
      public Supplement<Document>,
      public LocalFrameView::LifecycleNotificationObserver {
 public:
  static const char kSupplementName[];
  static AIPageContentAgent* From(Document&);
  static void BindReceiver(
      LocalFrame* frame,
      mojo::PendingReceiver<mojom::blink::AIPageContentAgent> receiver);

  static AIPageContentAgent* GetOrCreateForTesting(Document&);
#if DCHECK_IS_ON()
  // If enabled, the ContentNode tree will be automatically built on page load.
  static void EnableAutomaticActionableExtractionOnPageLoadForTesting(
      LocalFrame&);
#endif

  AIPageContentAgent(base::PassKey<AIPageContentAgent>, LocalFrame&);
  AIPageContentAgent(const AIPageContentAgent&) = delete;
  AIPageContentAgent& operator=(const AIPageContentAgent&) = delete;
  ~AIPageContentAgent() override;

  void Trace(Visitor* visitor) const override;

  // mojom::blink::AIPageContentAgent overrides.
  void GetAIPageContent(mojom::blink::AIPageContentOptionsPtr options,
                        GetAIPageContentCallback callback) override;
  void GetImageBytes(int32_t dom_node_id,
                     GetImageBytesCallback callback) override;

  // public for testing.
  mojom::blink::AIPageContentPtr GetAIPageContentInternal(
      const mojom::blink::AIPageContentOptions& options) const;
  // LocalFrameView::LifecycleNotificationObserver overrides.
  void DidFinishPostLifecycleSteps(const LocalFrameView&) override;

  String DumpContentNodeTreeForTest();
  String DumpContentNodeForTest(Node* node);

#if DCHECK_IS_ON()
  // Called by the DOMContentLoaded listener to kick off auto-build; this is
  // only used by tests to ensure we run cleanly without triggering crashes or
  // checks/dchecks.
  void RunAutoBuildAfterDOMContentLoadedForTesting();
#endif

 private:
  void GetAIPageContentSync(mojom::blink::AIPageContentOptionsPtr options,
                            GetAIPageContentCallback callback,
                            base::TimeTicks start_time) const;

  // Synchronously services a single request.
  class ContentBuilder {
    STACK_ALLOCATED();

   public:
    explicit ContentBuilder(const mojom::blink::AIPageContentOptions& options);
    ~ContentBuilder();

    mojom::blink::AIPageContentPtr Build(LocalFrame& frame);

   private:
    using NodeIdAttributeTypeAllowlist =
        base::EnumSet<mojom::blink::AIPageContentAttributeType,
                      mojom::blink::AIPageContentAttributeType::kRoot,
                      mojom::blink::AIPageContentAttributeType::kMaxValue>;

    class RecursionData {
      STACK_ALLOCATED();

     public:
      RecursionData(const ComputedStyle& document_style);

      bool is_aria_disabled = false;
      bool is_aria_hidden = false;
      bool is_in_fixed_pos_subtree = false;
      // The nearest overflow container clips descendants. It may or may not be
      // user-scrollable, because `overflow:hidden` also creates a container.
      const LayoutBox* nearest_overflow_container = nullptr;
      // Once an overflow container cannot be reached, its descendants cannot
      // be reached through that container either.
      bool is_inside_unreachable_overflow_container = false;
      const ComputedStyle& document_style;
      int stack_depth = 0;
      DOMNodeId accessibility_focused_node_id = kInvalidDOMNodeId;
    };

    bool actionable_mode() const {
      return options_->mode ==
             mojom::blink::AIPageContentMode::kActionableElements;
    }
    // Returns true if any descendant of `object` has a computed value of
    // visible for `visibility`. `ancestor_for_geometry_repair` is the nearest
    // APC ancestor that started with an empty outer box.
    bool WalkChildren(const LayoutObject& object,
                      mojom::blink::AIPageContentNode& content_node,
                      const RecursionData& recursion_data,
                      mojom::blink::AIPageContentNode*
                          ancestor_for_geometry_repair = nullptr);
    void ProcessIframe(const LayoutIFrame& object,
                       mojom::blink::AIPageContentNode& content_node,
                       const RecursionData& recursion_data);
    mojom::blink::AIPageContentNodePtr MaybeGenerateContentNode(
        const LayoutObject& object,
        const RecursionData& recursion_data);
    mojom::blink::AIPageContentNodePtr MaybeGenerateContentNodeImpl(
        const LayoutObject& object,
        const RecursionData& recursion_data);
    void AddPageInteractionInfo(const Document& document,
                                mojom::blink::AIPageContent& page_content);
    void AddFrameData(LocalFrame& frame,
                      mojom::blink::AIPageContentFrameData& frame_data);
    void AddFrameInteractionInfo(
        const LocalFrame& frame,
        mojom::blink::AIPageContentFrameInteractionInfo&
            frame_interaction_info);
    void MaybeAddPopupData(LocalFrame& frame,
                           mojom::blink::AIPageContentFrameData& frame_data);
    void AddNodeInteractionInfo(
        const LayoutObject& object,
        mojom::blink::AIPageContentAttributes& attributes,
        bool is_aria_disabled,
        bool is_aria_hidden);
    void AddInteractionInfoForHitTesting(
        const Node* node,
        mojom::blink::AIPageContentNodeInteractionInfo& interaction_info) const;
    void AddMetaData(
        const LocalFrame& frame,
        Vector<mojom::blink::AIPageContentMetaPtr>& meta_data) const;
    void AddNodeGeometry(const LayoutObject& object,
                         mojom::blink::AIPageContentAttributes& attributes,
                         DOMNodeId accessibility_focused_node_id);
    void AddAnnotatedRoles(const LayoutObject& object,
                           Vector<mojom::blink::AIPageContentAnnotatedRole>&
                               annotated_roles) const;
    void AddLabel(const LayoutObject& object,
                  mojom::blink::AIPageContentAttributes& attributes) const;
    // Adds the control node id if this is a label associated with a form
    // control. This includes both explicit association using for, or
    // implicit association when the input node is a descendant of the label
    // node.
    void PopulateLabelForDomNodeId(
        const LayoutObject& object,
        mojom::blink::AIPageContentAttributes& attributes);
    // Returns whether `attributes.dom_node_id` should be emitted in APC output.
    //
    // If `node_id_allowlist` is unset, emit all ids.
    //
    // Otherwise, emit ids for:
    // 1. Actionable targets (`node_interaction_info` in actionable mode).
    // 2. Metadata-linked nodes tracked in `interactive_dom_node_ids_` (focused
    //    element, accessibility focus, selection endpoints, label-for targets,
    //    popup openers).
    // 3. Attribute types listed in `node_id_allowlist`.
    // This method must not allocate a new DomNodeIds entry for nodes that are
    // ultimately suppressed by policy. Callers should only mint new ids after
    // this returns true.
    bool ShouldEmitNodeIdForOutput(
        const LayoutObject& object,
        const mojom::blink::AIPageContentAttributes& attributes) const;
    // Returns true when the caller policy allowlists this `attribute_type` for
    // id emission.
    bool IsNodeIdAttributeTypeAllowlisted(
        mojom::blink::AIPageContentAttributeType attribute_type) const;
    bool ShouldSkipSingleNode(
        const LayoutObject& object,
        const mojom::blink::AIPageContentAttributes& attributes) const;

    DOMNodeId AddInteractiveNode(Node& node);
    void AddInteractiveNode(DOMNodeId dom_node_id);
    void ComputeHitTestableNodesInViewport(const LocalFrame& frame);

    void UpdateLifecycle(Document& document);

    // Collects the visible bounding box for nodes that require redaction
    // (e.g. passwords or masked elements) so they can be obscured in
    // screenshots.
    void CollectGeometryForRedactedNodes(
        const LayoutObject& object,
        mojom::blink::AIPageContentRedactionDecision redaction_decision,
        std::optional<gfx::Rect> visible_bounding_box = std::nullopt);

    bool ShouldAddNodeGeometry(
        const LayoutObject& object,
        const mojom::blink::AIPageContentAttributes& attributes,
        DOMNodeId accessibility_focused_node_id) const;

    Vector<gfx::Rect> visible_bounding_boxes_for_redaction_;

    // The set of node ids that must always be emitted in APC output for
    // round-trippable metadata and interaction flows.
    HashSet<DOMNodeId, IntWithZeroKeyHashTraits<DOMNodeId>>
        interactive_dom_node_ids_;

    const raw_ref<const mojom::blink::AIPageContentOptions> options_;

    // Keyed by Node identity within a single extraction pass.
    HeapHashMap<Member<const Node>, int32_t> dom_node_to_z_order_;

    // Whether the stack depth has exceeded the max tree depth.
    bool stack_depth_exceeded_ = false;

    // List of nodes marked as isAccessibleForFree=false.
    PaidContent paid_content_;

    // Built once per extraction so node-id policy checks do not rescan the
    // allowlisted attribute types for every content node in the tree.
    NodeIdAttributeTypeAllowlist allowlisted_attribute_types_;
  };

  void Bind(mojo::PendingReceiver<mojom::blink::AIPageContentAgent> receiver);
  void EnsureLifecycleObserverRegistered();

  HeapMojoReceiverSet<mojom::blink::AIPageContentAgent, AIPageContentAgent>
      receiver_set_;
  // Already registered for lifetime notifications.
  bool is_lifecycle_observer_registered_ = false;
  // Tasks to run when post lifecycle.
  Vector<base::OnceClosure> async_extraction_tasks_;

#if DCHECK_IS_ON()
  void ListenForDOMContentLoadedForAutoBuild();
  AutoBuildHelper* GetOrCreateAutoBuildHelper();
  Member<AutoBuildHelper> auto_build_helper_;
  friend class AutoBuildHelper;
#endif
};

}  // namespace blink

#endif  // THIRD_PARTY_BLINK_RENDERER_MODULES_CONTENT_EXTRACTION_AI_PAGE_CONTENT_AGENT_H_