// Copyright 2018 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_LOAD_METRICS_BROWSER_PAGE_LOAD_METRICS_TEST_WAITER_H_
#define COMPONENTS_PAGE_LOAD_METRICS_BROWSER_PAGE_LOAD_METRICS_TEST_WAITER_H_

#include <memory>
#include <unordered_set>
#include <vector>

#include "base/byte_size.h"
#include "base/memory/weak_ptr.h"
#include "base/run_loop.h"
#include "base/time/time.h"
#include "components/page_load_metrics/browser/metrics_lifecycle_observer.h"
#include "components/page_load_metrics/browser/page_load_tracker.h"
#include "components/page_load_metrics/common/page_load_metrics.mojom.h"
#include "content/public/browser/render_frame_host.h"
#include "third_party/blink/public/common/use_counter/use_counter_feature_tracker.h"
#include "ui/gfx/geometry/size.h"

namespace page_load_metrics {

class WaiterMetricsObserver;

class PageLoadMetricsTestWaiter : public MetricsLifecycleObserver {
 public:
  // A bitvector to express which timing fields to match on.
  enum class TimingField : int {
    kFirstPaint = 1 << 0,
    kFirstContentfulPaint = 1 << 1,
    kFirstMeaningfulPaint = 1 << 2,
    // kDocumentWriteBlockReload is deprecated.
    kDocumentWriteBlockReload = 1 << 3,
    kLoadEvent = 1 << 4,
    // kLoadTimingInfo waits for main frame timing info only.
    kLoadTimingInfo = 1 << 5,
    kLargestContentfulPaint = 1 << 6,
    kFirstInputOrScroll = 1 << 7,
    kFirstInputDelay = 1 << 8,
    kFirstPaintAfterBackForwardCacheRestore = 1 << 9,
    kFirstInputDelayAfterBackForwardCacheRestore = 1 << 10,
    kRequestAnimationFrameAfterBackForwardCacheRestore = 1 << 11,
    kFirstScrollDelay = 1 << 12,
    kSoftNavigationCountUpdated = 1 << 13,
    kMonotonicFirstPaint = 1 << 14,
    kMonotonicFirstContentfulPaint = 1 << 15,
  };

  // Identify which frame the layout shift happens.
  enum class ShiftFrame {
    LayoutShiftOnlyInMainFrame,
    LayoutShiftOnlyInSubFrame,
    LayoutShiftOnlyInBothFrames,
    NoLayoutShift,
  };

  explicit PageLoadMetricsTestWaiter(content::WebContents* web_contents);
  PageLoadMetricsTestWaiter(content::WebContents* web_contents,
                            const char* observer_name);

  ~PageLoadMetricsTestWaiter() override;

  // Add a page-level expectation.
  void AddPageExpectation(TimingField field);

  // Add a subframe-level expectation.
  void AddSubFrameExpectation(TimingField field);

  // Add a page-level expectation on the k-th back-forward cache restoration.
  // Note:
  // * |back_forward_timings_index| is 0-based.
  // * Only k[x]AfterBackForwardCacheRestore are valid arguments for |field|,
  //   with [x] one of {FirstPaint, FirstInput, RequestAnimationFrame}.
  void AddPageBackForwardCacheRestoreExpectation(
      size_t back_forward_timings_index,
      TimingField field);

  // Add a frame size expectation. Expects that at least one frame receives a
  // size update of |size|.
  void AddFrameSizeExpectation(const gfx::Size& size);

  // Add a main frame rectangle expectation. Expects that a frame receives an
  // update with main frame rectangle `rect`. Subsequent calls overwrite unmet
  // expectations.
  void AddMainFrameRectExpectation(const gfx::Rect& rect);

  // Indicates that we expect at least one main frame rectangle update, with any
  // rect allowed.
  // TODO(skobes): Unify this API with AddMainFrameRectExpectation.
  void SetMainFrameRectExpectation();

  // Indicates that we expect at least one notification for the main frame ad
  // rectangles update, with any rect allowed.
  void SetMainFrameAdRectsExpectation();

  // Add a main frame viewport rectangle expectation. Expects that the main
  // frame receives its viewport rectangle in the main frame's coordinate
  // system. Subsequent calls overwrite unmet expectations.
  void AddMainFrameViewportRectExpectation(const gfx::Rect& rect);

  // Add a single WebFeature expectation.
  void AddWebFeatureExpectation(blink::mojom::WebFeature web_feature);

  // Add a single UseCounterFeature expectation.
  void AddUseCounterFeatureExpectation(const blink::UseCounterFeature& feature);

  // Add a custom user timing mark expectation.
  void AddCustomUserTimingMarkExpectation(const std::string& mark_name);

  // Wait for the subframe to navigate at least once.
  void AddSubframeNavigationExpectation();

  // Wait for the subframe to load at least one byte.
  void AddSubframeDataExpectation();

  // Add a minimum completed resource expectation.
  void AddMinimumCompleteResourcesExpectation(
      int expected_minimum_complete_resources);

  // Add aggregate received resource bytes expectation.
  void AddMinimumNetworkBytesExpectation(
      base::ByteSize expected_minimum_network_bytes);

  // Add aggregate time spent in cpu for page expectation.
  void AddMinimumAggregateCpuTimeExpectation(base::TimeDelta minimum);

  // Adds all |blink::LoadingBehaviorFlag|s set in |behavior_flags| to the
  // set of expected behaviors.
  void AddLoadingBehaviorExpectation(int behavior_flags);

  // Add minimum largest contentful paint image update count to be expected.
  // Also reset observed largest contentful paint image count.
  void AddMinimumLargestContentfulPaintImageExpectation(int expected_minumum);

  // Add minimum largest contentful paint text update count to be expected.
  // Also reset observed largest contentful paint text count.
  void AddMinimumLargestContentfulPaintTextExpectation(int expected_minumum);

  void AddLargestContentfulPaintGreaterThanExpectation(double timestamp);

  void AddSoftNavigationCountExpectation(int expected_count);

  void AddSoftNavigationLargestContentfulPaintExpectation(int expected_count);

  // Add a main/sub frame layout shift expectation.
  void AddPageLayoutShiftExpectation(
      ShiftFrame frame = ShiftFrame::LayoutShiftOnlyInMainFrame,
      uint64_t num_layout_shifts = 1);

  // Adds a condition to wait for OnComplete invocation that indicates the
  // observer will be gone, and Wait() can ensure all metrics are recorded.
  void AddOnCompleteCalledExpectation();

  // Whether the given TimingField was observed in the page.
  bool DidObserveInPage(TimingField field) const;

  // Whether the given WebFeature was observed in the page.
  bool DidObserveWebFeature(blink::mojom::WebFeature feature) const;

  // Whether the given ad rect was observed in the page.
  bool DidObserveMainFrameAdRect(const gfx::Rect& rect) const;

  // Waits for PageLoadMetrics events that match the fields set by the add
  // expectation methods. All matching fields must be set to end this wait.
  // All expectations are reset when the wait ends.
  void Wait();

  base::ByteSize current_network_bytes() const {
    return current_network_bytes_;
  }

  base::ByteSize current_network_body_bytes() const {
    return current_network_body_bytes_;
  }

  // Add the number of interactions count expectation.
  void AddNumInteractionsExpectation(uint64_t expected_num_interactions) {
    expected_num_interactions_ = expected_num_interactions;
  }

 protected:
  virtual bool ExpectationsSatisfied() const;
  void AssertExpectationsSatisfied() const;

  // Intended to be overridden in tests to allow tests to wait on other resource
  // conditions.
  virtual void HandleResourceUpdate(
      const page_load_metrics::mojom::ResourceDataUpdatePtr& resource) {}

  // Resets all expectations.
  virtual void ResetExpectations();

 private:
  const char* observer_name_;

  // Manages a bitset of TimingFields.
  class TimingFieldBitSet {
   public:
    TimingFieldBitSet() = default;

    // Returns whether this bitset has all bits unset.
    bool Empty() const { return bitmask_ == 0; }

    // Returns whether this bitset has the given bit set.
    bool IsSet(TimingField field) const {
      return (bitmask_ & static_cast<int>(field)) != 0;
    }

    // Sets the bit for the given |field|.
    void Set(TimingField field) { bitmask_ |= static_cast<int>(field); }

    // Clears the bit for the given |field|.
    void Clear(TimingField field) { bitmask_ &= ~static_cast<int>(field); }

    // Merges bits set in |other| into this bitset.
    void Merge(const TimingFieldBitSet& other) { bitmask_ |= other.bitmask_; }

    // Clears all bits set in the |other| bitset.
    void ClearMatching(const TimingFieldBitSet& other) {
      bitmask_ &= ~other.bitmask_;
    }

    // Returns whether all the bits in this bitset are set in |other|.
    bool AreAllSetIn(const TimingFieldBitSet& other) const {
      return !((bitmask_ & other.bitmask_) ^ bitmask_);
    }

    // Returns the string representation of the TimingFields this bitset
    // contains. This method is not called anywhere and is for debug purpose
    // only.
    std::string ToDebugString() const;

    // Returns true if the bitset contains the TimingField. This method is not
    // called anywhere and is for debug purpose only.
    bool ContainsTimingField(TimingField time_field) const {
      return (bitmask_ & static_cast<int>(time_field)) > 0;
    }

   private:
    int bitmask_ = 0;
  };

  struct FrameSizeComparator {
    bool operator()(const gfx::Size a, const gfx::Size b) const;
  };

  TimingFieldBitSet GetMatchedBits(
      const page_load_metrics::mojom::PageLoadTiming& timing,
      const page_load_metrics::mojom::FrameMetadata& metadata);

  // Updates observed page fields when a timing update is received by the
  // MetricsWebContentsObserver. Stops waiting if expectations are satsfied
  // after update.
  void OnTimingUpdated(content::RenderFrameHost* subframe_rfh,
                       const page_load_metrics::mojom::PageLoadTiming& timing);

  // Invoked when a soft navigation is detected by the
  // PageLoadMetricsUpdateDispatcher.
  void OnSoftNavigation();

  // Invoked when one or more soft LCPs are detected by the
  // PageLoadMetricsUpdateDispatcher.
  void OnSoftNavigationLargestContentfulPaint(uint64_t num_soft_lcps);

  // Updates observed page fields when a input timing update is received by the
  // MetricsWebContentsObserver. Stops waiting if expectations are satsfied
  // after update.
  void OnPageEventTimingUpdated(uint64_t num_interactions);

  // Updates observed page fields when a timing update is received by the
  // MetricsWebContentsObserver. Stops waiting if expectations are satsfied
  // after update.
  void OnCpuTimingUpdated(content::RenderFrameHost* subframe_rfh,
                          const page_load_metrics::mojom::CpuTiming& timing);

  // Updates observed page fields when a loading behavior (see
  // |blink::LoadingBehaviorFlag|) is observed by MetricsWebContentsObserver.
  // Stops waiting if expectations are satsfied after update.
  void OnLoadingBehaviorObserved(int behavior_flags);

  // Updates observed page fields when a resource load is observed by
  // MetricsWebContentsObserver.  Stops waiting if expectations are satsfied
  // after update.
  void OnLoadedResource(const page_load_metrics::ExtraRequestCompleteInfo&
                            extra_request_complete_info);

  // Updates counters as updates are received from a resource load. Stops
  // waiting if expectations are satisfied after update.
  void OnResourceDataUseObserved(
      content::RenderFrameHost* rfh,
      const std::vector<page_load_metrics::mojom::ResourceDataUpdatePtr>&
          resources);

  // Updates |observed_.web_features_| to record any new feature observed.
  // Stops waiting if expectations are satisfied after update.
  void OnFeaturesUsageObserved(
      content::RenderFrameHost* rfh,
      const std::vector<blink::UseCounterFeature>& features);

  void OnCustomUserTimingMarksObserved(
      content::RenderFrameHost* rfh,
      const std::vector<page_load_metrics::mojom::CustomUserTimingMarkPtr>&
          timings);

  // Updates |observed_.layout_shift_| to record any update of new layout
  // shift. Stops waiting if expectations are satisfied after update.
  void OnPageRenderDataUpdate(const mojom::FrameRenderDataUpdate& render_data,
                              bool is_main_frame);

  void FrameSizeChanged(content::RenderFrameHost* render_frame_host,
                        const gfx::Size& frame_size);

  void OnMainFrameRectChanged(const gfx::Rect& main_frame_rect);

  void OnMainFrameViewportRectChanged(
      const gfx::Rect& main_frame_viewport_rect);

  void OnMainFrameAdRectsChanged(
      const base::flat_map<int, gfx::Rect>& main_frame_ad_rects);

  void OnDidFinishSubFrameNavigation(
      content::NavigationHandle* navigation_handle);

  void OnComplete(const mojom::PageLoadTiming& timing);

  void OnTrackerCreated(page_load_metrics::PageLoadTracker* tracker) override;

  void OnCommit(page_load_metrics::PageLoadTracker* tracker) override;

  void OnActivate(page_load_metrics::PageLoadTracker* tracker) override;

  // These methods check whether expectations are satisfied for specific fields
  // inside the State object, by comparing them in expected_ and observed_.
  bool CpuTimeExpectationsSatisfied() const;
  bool LoadingBehaviorExpectationsSatisfied() const;
  bool ResourceUseExpectationsSatisfied() const;
  bool UseCounterExpectationsSatisfied() const;
  bool SubframeNavigationExpectationsSatisfied() const;
  bool SubframeDataExpectationsSatisfied() const;
  bool MainFrameRectExpectationsSatisfied() const;
  bool MainFrameViewportRectExpectationsSatisfied() const;
  bool MainFrameAdRectsExpectationsSatisfied() const;
  bool LayoutShiftExpectationsSatisfied() const;
  bool NumInteractionsExpectationsSatisfied() const;
  bool NumLargestContentfulPaintImageSatisfied() const;
  bool NumLargestContentfulPaintTextSatisfied() const;
  bool LargestContentfulPaintGreaterThanExpectationSatisfied() const;
  bool SoftNavigationCountExpectationSatisfied() const;
  bool SoftNavigationLargestContentfulPaintExpectationSatisfied() const;
  bool CustomUserTimingMarksExpectationsSatisfied() const;

  void AddObserver(page_load_metrics::PageLoadTracker* tracker);

  std::unique_ptr<base::RunLoop> run_loop_;

  // Holds information about events that can be expected or observed. Each call
  // to Wait() compares expected_ to observed_, and resets both.
  struct State {
    State();
    ~State();

    TimingFieldBitSet page_fields_;
    TimingFieldBitSet subframe_fields_;
    base::flat_map<size_t, TimingFieldBitSet> page_bfcache_restore_fields_;
    blink::UseCounterFeatureTracker feature_tracker_;
    std::unordered_set<std::string> custom_user_timing_marks_;
    int loading_behavior_flags_ = 0;
    bool subframe_navigation_ = false;
    bool subframe_data_ = false;
    std::set<gfx::Size, FrameSizeComparator> frame_sizes_;
    bool did_set_main_frame_rect_ = false;
    bool did_observed_main_frame_ad_rects_ = false;
    std::optional<gfx::Rect> main_frame_rect_;
    std::optional<gfx::Rect> main_frame_viewport_rect_;
    uint64_t num_layout_shifts_ = 0;
    bool on_complete_ = false;
  };
  State expected_;
  State observed_;

  int current_complete_resources_ = 0;
  base::ByteSize current_network_bytes_;

  // The last observed main frame ad rectangle for each id. This doesn't
  // get reset in `ResetExpectations`.
  base::flat_map<int, gfx::Rect> main_frame_ad_rects_;

  // Network body bytes are only counted for complete resources.
  base::ByteSize current_network_body_bytes_;
  int expected_minimum_complete_resources_ = 0;
  base::ByteSize expected_minimum_network_bytes_;

  // Total time spent int the cpu aggregated across the frames on the page.
  base::TimeDelta current_aggregate_cpu_time_;
  base::TimeDelta expected_minimum_aggregate_cpu_time_;

  bool attach_on_tracker_creation_ = false;
  bool did_add_observer_ = false;
  bool soft_navigation_count_updated_ = false;

  uint64_t current_num_interactions_ = 0;
  uint64_t expected_num_interactions_ = 0;

  uint64_t expected_num_largest_contentful_paint_image_ = 0;
  uint64_t current_num_largest_contentful_paint_image_ = 0;

  uint64_t expected_num_largest_contentful_paint_text_ = 0;
  uint64_t current_num_largest_contentful_paint_text_ = 0;

  uint64_t expected_soft_navigation_count_ = 0;
  uint64_t current_soft_navigation_count_ = 0;

  uint64_t expected_num_soft_navigation_largest_contentful_paint_ = 0;
  uint64_t current_num_soft_navigation_largest_contentful_paint_ = 0;

  double expected_min_largest_contentful_paint_ = -1.0;
  double observed_largest_contentful_paint_ = 0.0;

  ShiftFrame shift_frame_ = ShiftFrame::NoLayoutShift;

  base::WeakPtrFactory<PageLoadMetricsTestWaiter> weak_factory_{this};

  friend class WaiterMetricsObserver;
};

}  // namespace page_load_metrics

#endif  // COMPONENTS_PAGE_LOAD_METRICS_BROWSER_PAGE_LOAD_METRICS_TEST_WAITER_H_