// 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.

#import "ios/chrome/browser/web/model/image_fetch/image_fetch_tab_helper.h"

#import "base/base64.h"
#import "base/functional/bind.h"
#import "base/json/string_escape.h"
#import "base/metrics/histogram_macros.h"
#import "base/strings/utf_string_conversions.h"
#import "base/values.h"
#import "components/image_fetcher/core/image_data_fetcher.h"
#import "ios/chrome/browser/shared/model/profile/profile_ios.h"
#import "ios/chrome/browser/web/model/image_fetch/image_fetch_java_script_feature.h"
#import "ios/web/common/referrer_util.h"
#import "ios/web/public/js_messaging/web_frame.h"
#import "ios/web/public/js_messaging/web_frames_manager.h"
#import "ios/web/public/navigation/navigation_context.h"
#import "ios/web/public/thread/web_task_traits.h"
#import "ios/web/public/thread/web_thread.h"
#import "services/network/public/cpp/shared_url_loader_factory.h"

const char kUmaGetImageDataByJsResult[] =
    "Mobile.ContextMenu.GetImageDataByJsResult";

namespace {
// Key for image_fetcher
const char kImageFetcherKeyName[] = "0";
// Timeout for GetImageDataByJs.
constexpr base::TimeDelta kGetImageDataByJsTimeout = base::Milliseconds(300);

// Wrapper class for image_fetcher::IOSImageDataFetcherWrapper. ImageFetcher is
// attached to ProfileIOS instead of web::WebState, because if a user closes the
// tab immediately after Copy/Save image, the web::WebState will be destroyed
// thus fail the download.
class ImageFetcher : public image_fetcher::ImageDataFetcher,
                     public base::SupportsUserData::Data {
 public:
  ImageFetcher(const ImageFetcher&) = delete;
  ImageFetcher& operator=(const ImageFetcher&) = delete;

  ~ImageFetcher() override = default;

  ImageFetcher(
      scoped_refptr<network::SharedURLLoaderFactory> url_loader_factory)
      : image_fetcher::ImageDataFetcher(url_loader_factory) {}

  // Helper that retrieves the ImageFetcher attached to the WebState's profile.
  static ImageFetcher* FromWebState(web::WebState* web_state) {
    return FromProfile(
        ProfileIOS::FromBrowserState(web_state->GetBrowserState()));
  }

  // Retrieves the ImageFetcher for `profile` creating it first if it does
  // not exist yet.
  static ImageFetcher* FromProfile(ProfileIOS* profile) {
    if (ImageFetcher* image_fetcher = static_cast<ImageFetcher*>(
            profile->GetUserData(&kImageFetcherKeyName))) {
      return image_fetcher;
    }

    profile->SetUserData(
        &kImageFetcherKeyName,
        std::make_unique<ImageFetcher>(profile->GetSharedURLLoaderFactory()));
    return static_cast<ImageFetcher*>(
        profile->GetUserData(&kImageFetcherKeyName));
  }
};
}  // namespace

ImageFetchTabHelper::ImageFetchTabHelper(web::WebState* web_state)
    : web_state_(web_state), weak_ptr_factory_(this) {
  web_state->AddObserver(this);
}

ImageFetchTabHelper::~ImageFetchTabHelper() = default;

void ImageFetchTabHelper::DidStartNavigation(
    web::WebState* web_state,
    web::NavigationContext* navigation_context) {
  if (navigation_context->IsSameDocument()) {
    return;
  }
  for (auto&& pair : js_callbacks_) {
    std::move(pair.second).Run(nullptr);
  }
  js_callbacks_.clear();
}

void ImageFetchTabHelper::WebStateDestroyed(web::WebState* web_state) {
  for (auto&& pair : js_callbacks_) {
    std::move(pair.second).Run(nullptr);
  }
  web_state->RemoveObserver(this);
  web_state_ = nullptr;
}

void ImageFetchTabHelper::GetImageData(const GURL& url,
                                       const web::Referrer& referrer,
                                       const std::string& frame_id,
                                       const url::Origin& frame_origin,
                                       ImageDataCallback callback) {
  web::WebFrame* frame = ImageFetchJavaScriptFeature::GetInstance()
                             ->GetWebFramesManager(web_state_)
                             ->GetFrameWithId(frame_id);
  if (!frame) {
    FetchImageDataWithFetcher(url, referrer, callback);
    return;
  }

  if (frame->GetSecurityOrigin() != frame_origin) {
    FetchImageDataWithFetcher(url, referrer, callback);
    return;
  }

  // `this` is captured into the callback of GetImageDataByJs, which will always
  // be invoked before the `this` is destroyed, so it's safe.
  GetImageDataByJs(
      url, frame, kGetImageDataByJsTimeout,
      base::BindOnce(&ImageFetchTabHelper::JsCallbackOfGetImageData,
                     base::Unretained(this), url, referrer, callback));
}

void ImageFetchTabHelper::FetchImageDataWithFetcher(
    const GURL& url,
    const web::Referrer& referrer,
    ImageDataCallback callback) {
  ImageFetcher::FromWebState(web_state_)
      ->FetchImageData(
          url,
          base::BindOnce(^(const std::string& image_data,
                           const image_fetcher::RequestMetadata& metadata) {
            NSData* nsdata = [NSData dataWithBytes:image_data.data()
                                            length:image_data.size()];
            callback(nsdata);
          }),
          web::ReferrerHeaderValueForNavigation(url, referrer),
          web::PolicyForNavigation(url, referrer), NO_TRAFFIC_ANNOTATION_YET,
          /*send_cookies=*/true);
}

void ImageFetchTabHelper::JsCallbackOfGetImageData(
    const GURL& url,
    const web::Referrer& referrer,
    ImageDataCallback callback,
    const std::string* data) {
  if (data) {
    callback([NSData dataWithBytes:data->c_str() length:data->size()]);
    return;
  }
  FetchImageDataWithFetcher(url, referrer, callback);
}

void ImageFetchTabHelper::GetImageDataByJs(const GURL& url,
                                           web::WebFrame* frame,
                                           base::TimeDelta timeout,
                                           JsCallback&& callback) {
  if (!frame) {
    std::move(callback).Run(nullptr);
    return;
  }

  ++call_id_;
  DCHECK_EQ(js_callbacks_.count(call_id_), 0UL);
  js_callbacks_.insert({call_id_, std::move(callback)});

  web::GetUIThreadTaskRunner({})->PostDelayedTask(
      FROM_HERE,
      base::BindRepeating(&ImageFetchTabHelper::OnJsTimeout,
                          weak_ptr_factory_.GetWeakPtr(), call_id_),
      timeout);

  ImageFetchJavaScriptFeature::GetInstance()->GetImageData(frame, call_id_,
                                                           url);
}

void ImageFetchTabHelper::RecordGetImageDataByJsResult(
    ContextMenuGetImageDataByJsResult result) {
  UMA_HISTOGRAM_ENUMERATION(kUmaGetImageDataByJsResult, result);
}

void ImageFetchTabHelper::HandleJsSuccess(int call_id,
                                          std::string& decoded_data,
                                          std::string& from) {
  if (!js_callbacks_.count(call_id)) {
    return;
  }
  JsCallback callback = std::move(js_callbacks_[call_id]);
  js_callbacks_.erase(call_id);

  DCHECK(!decoded_data.empty());
  std::move(callback).Run(&decoded_data);

  if (from == "canvas") {
    RecordGetImageDataByJsResult(
        ContextMenuGetImageDataByJsResult::kCanvasSucceed);
  } else if (from == "xhr") {
    RecordGetImageDataByJsResult(
        ContextMenuGetImageDataByJsResult::kXMLHttpRequestSucceed);
  }
}

void ImageFetchTabHelper::HandleJsFailure(int call_id) {
  if (!js_callbacks_.count(call_id)) {
    return;
  }
  JsCallback callback = std::move(js_callbacks_[call_id]);
  js_callbacks_.erase(call_id);

  std::move(callback).Run(nullptr);
  RecordGetImageDataByJsResult(ContextMenuGetImageDataByJsResult::kFail);
}

void ImageFetchTabHelper::OnJsTimeout(int call_id) {
  if (js_callbacks_.count(call_id)) {
    JsCallback callback = std::move(js_callbacks_[call_id]);
    js_callbacks_.erase(call_id);
    std::move(callback).Run(nullptr);
    RecordGetImageDataByJsResult(ContextMenuGetImageDataByJsResult::kTimeout);
  }
}
