// Copyright 2023 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/snapshots/model/legacy_snapshot_generator.h"

#import "base/debug/crash_logging.h"
#import "base/debug/dump_without_crashing.h"
#import "base/functional/bind.h"
#import "build/blink_buildflags.h"
#import "ios/chrome/browser/shared/model/url/chrome_url_constants.h"
#import "ios/chrome/browser/snapshots/model/model_swift.h"
#import "ios/chrome/browser/snapshots/model/snapshot_scale.h"
#import "ios/chrome/browser/snapshots/model/snapshot_source_tab_helper.h"
#import "ios/chrome/browser/snapshots/model/web_state_snapshot_info.h"
#import "ios/web/public/thread/web_thread.h"
#import "ios/web/public/web_client.h"
#import "ios/web/public/web_state.h"

namespace {

// Contains information needed for snapshotting.
struct SnapshotInfo {
  UIView* baseView;
  CGRect snapshotFrameInBaseView;
  CGRect snapshotFrameInWindow;
};

}  // namespace

@implementation LegacySnapshotGenerator {
  // The associated WebState.
  base::WeakPtr<web::WebState> _webState;
}

- (instancetype)initWithWebState:(web::WebState*)webState {
  if ((self = [super init])) {
    DCHECK(webState);
    _webState = webState->GetWeakPtr();
  }
  return self;
}

- (void)generateSnapshotWithCompletion:(void (^)(UIImage*))completion {
  [self generateSnapshotWithCompletion:completion includeOverlays:YES];
}

- (void)generateSnapshotWithoutOverlaysWithCompletion:
    (void (^)(UIImage*))completion {
  [self generateSnapshotWithCompletion:completion includeOverlays:NO];
}

- (UIImage*)generateUIViewSnapshot {
  if (![self canTakeSnapshot] || !_webState) {
    return nil;
  }
  [_delegate
      willUpdateSnapshotWithWebStateInfo:[[WebStateSnapshotInfo alloc]
                                             initWithWebState:_webState.get()]];

  std::optional<SnapshotInfo> snapshotInfo = [self snapshotInfo];
  if (!snapshotInfo) {
    return nil;
  }
  // Ideally, generate an UIImage by one step with `UIGraphicsImageRenderer`,
  // however, it generates a black image when the size of `baseView` is larger
  // than `frameInBaseView`. So this is a workaround to generate an UIImage by
  // dividing the step into 2 steps; 1) convert an UIView to an UIImage 2) crop
  // an UIImage with `frameInBaseView`.
  UIImage* baseImage = [self convertFromBaseView:snapshotInfo.value().baseView];
  return [self cropImage:baseImage
         frameInBaseView:snapshotInfo.value().snapshotFrameInBaseView];
}

- (UIImage*)generateUIViewSnapshotWithOverlays {
  if (![self canTakeSnapshot]) {
    return nil;
  }
  std::optional<SnapshotInfo> snapshotInfo = [self snapshotInfo];
  if (!snapshotInfo) {
    return nil;
  }
  return [self addOverlays:[self overlays]
                 baseImage:[self generateUIViewSnapshot]
             frameInWindow:snapshotInfo.value().snapshotFrameInWindow];
}

#pragma mark - Private methods

// Generates a new snapshot and runs a callback with the new snapshot image.
// The generated image includes overlays if `includeOverlays` is YES.
// - If the web state is not showing a new tab page, the page is not incognito
//   and it doesn't have JavaScript dialogs,
//   - it uses WebKit-based snapshot API
//   - and the callback is called asynchronously.
// - Otherwise,
//   - it uses UIKit-based snapshot API
//   - and the callback is called immediately (without posting a task).
- (void)generateSnapshotWithCompletion:(void (^)(UIImage*))completion
                       includeOverlays:(BOOL)includeOverlays {
  // TODO(crbug.com/452299163): A last committed URL is nil when this method is
  // called before the navigation has been committed. For instance, a snapshot
  // is taken prior to the navigation when opening multiple new tab pages
  // quickly. The execution continues but it may choose the wrong way to take a
  // snapshot (-generateWKWebViewSnapshotWithCompletion: vs
  // -generateUIViewSnapshotWithOverlays:).
  bool isNTP = _webState->GetLastCommittedURL() == kChromeUINewTabURL;
  SnapshotSourceTabHelper* snapshotSource =
      SnapshotSourceTabHelper::FromWebState(_webState.get());
  if (!isNTP && snapshotSource->CanTakeSnapshot()) {
    // Take the snapshot using the optimized WKWebView snapshotting API for
    // pages loaded in the web view when the WebState snapshot API is available.
    [self generateWKWebViewSnapshotWithCompletion:completion
                                  includeOverlays:includeOverlays];
    return;
  }
  // Use the UIKit-based snapshot API as a fallback when the WKWebView API is
  // unavailable.
  UIImage* snapshot = includeOverlays
                          ? [self generateUIViewSnapshotWithOverlays]
                          : [self generateUIViewSnapshot];
  if (completion) {
    completion(snapshot);
  }
}

// Asynchronously generates a new snapshot with WebKit-based snapshot API and
// runs a callback with the new snapshot image. It is an error to call this
// method if the web state is showing anything other (e.g., native content) than
// a web view.
- (void)generateWKWebViewSnapshotWithCompletion:(void (^)(UIImage*))completion
                                includeOverlays:(BOOL)includeOverlays {
  if (![self canTakeSnapshot]) {
    if (completion) {
      // Post a task to the current thread (UI thread).
      base::SequencedTaskRunner::GetCurrentDefault()->PostTask(
          FROM_HERE, base::BindOnce(completion, nil));
    }
    return;
  }
  SnapshotSourceTabHelper* snapshotSource =
      SnapshotSourceTabHelper::FromWebState(_webState.get());
  CHECK(snapshotSource->CanTakeSnapshot());

  [_delegate
      willUpdateSnapshotWithWebStateInfo:[[WebStateSnapshotInfo alloc]
                                             initWithWebState:_webState.get()]];

  std::optional<SnapshotInfo> snapshotInfo = [self snapshotInfo];
  if (!snapshotInfo) {
    if (completion) {
      // Post a task to the current thread (UI thread).
      base::SequencedTaskRunner::GetCurrentDefault()->PostTask(
          FROM_HERE, base::BindOnce(completion, nil));
    }
    return;
  }

  auto wrappedCompletion =
      ^(__weak LegacySnapshotGenerator* generator, UIImage* image) {
        UIImage* snapshot =
            [generator adjustWKWebViewSnapshotIfNecessary:image
                                          includeOverlays:includeOverlays];
        if (completion) {
          completion(snapshot);
        }
      };

  __weak LegacySnapshotGenerator* weakSelf = self;
  snapshotSource->TakeSnapshot(
      snapshotInfo.value().snapshotFrameInBaseView,
      base::BindRepeating(wrappedCompletion, weakSelf));
}

// Adjusts a snapshot taken by WebKit API if necessary.
// If the image is smaller than the base view, we need to add a background to
// the image (e.g. 1 page PDF in WKWebView. See crbug.com/399702753). Add
// overlays as well if they exist and `includeOverlays` is YES.
- (UIImage*)adjustWKWebViewSnapshotIfNecessary:(UIImage*)image
                               includeOverlays:(BOOL)includeOverlays {
  std::optional<SnapshotInfo> snapshotInfo = [self snapshotInfo];
  if (!snapshotInfo) {
    return nil;
  }
  CGRect frameInBaseView = snapshotInfo.value().snapshotFrameInBaseView;

  // If the image generated by WebKit API is smaller than originally
  // demanded, combine it with the background image.
  if (image.size.height < frameInBaseView.size.height) {
    UIImage* backgroundImage = [self generateUIViewSnapshot];

    UIGraphicsImageRendererFormat* format =
        [UIGraphicsImageRendererFormat preferredFormat];
    format.scale = [SnapshotImageScale floatImageScaleForDevice];
    format.opaque = YES;

    UIGraphicsImageRenderer* renderer =
        [[UIGraphicsImageRenderer alloc] initWithSize:frameInBaseView.size
                                               format:format];
    image = [renderer imageWithActions:^(
                          UIGraphicsImageRendererContext* UIContext) {
      [backgroundImage drawInRect:(CGRect){.origin = CGPointZero,
                                           .size = backgroundImage.size}];
      [image drawInRect:(CGRect){.origin = CGPointZero, .size = image.size}];
    }];
  }

  if (includeOverlays) {
    return [self addOverlays:[self overlays]
                   baseImage:image
               frameInWindow:snapshotInfo.value().snapshotFrameInWindow];
  }
  return image;
}

// Returns NO if WebState or the view is not ready for snapshot.
- (BOOL)canTakeSnapshot {
  // This allows for easier unit testing of classes that use SnapshotGenerator.
  if (!_delegate || !_webState) {
    return NO;
  }

  // Do not generate a snapshot if web usage is disabled (as the WebState's
  // view is blank in that case).
  if (!_webState->IsWebUsageEnabled()) {
    return NO;
  }

  return [_delegate
      canTakeSnapshotWithWebStateInfo:[[WebStateSnapshotInfo alloc]
                                          initWithWebState:_webState.get()]];
}

// Converts an UIView to an UIImage. The size of generated UIImage is the same
// as `baseView`.
- (UIImage*)convertFromBaseView:(UIView*)baseView {
  DCHECK(baseView);

  // Disable the automatic view dimming UIKit performs if a view is presented
  // modally over `baseView`.
  baseView.tintAdjustmentMode = UIViewTintAdjustmentModeNormal;

  // Note: When not using device scale, the output image size may slightly
  // differ from the input size due to rounding.
  const CGFloat kScale = [SnapshotImageScale floatImageScaleForDevice];
  DCHECK_GE(kScale, 1.0);
  UIGraphicsImageRendererFormat* format =
      [UIGraphicsImageRendererFormat preferredFormat];
  format.scale = kScale;
  format.opaque = NO;

  UIGraphicsImageRenderer* renderer =
      [[UIGraphicsImageRenderer alloc] initWithBounds:baseView.bounds
                                               format:format];

  __block BOOL snapshotSuccess = YES;
  UIImage* image =
      [renderer imageWithActions:^(UIGraphicsImageRendererContext* UIContext) {
        if (@available(iOS 26, *)) {
          // A translucent background starting from iOS 26 only works well with
          // drawViewHierarchyInRect.
          snapshotSuccess = [baseView drawViewHierarchyInRect:baseView.bounds
                                           afterScreenUpdates:NO];
        } else {
          // Take animations into account by rendering the presentation layer.
          // Fallback to the rendering the layer if not possible.
          CALayer* layerToRender =
              baseView.layer.presentationLayer ?: baseView.layer;
          // To mitigate against crashes like crbug.com/1429512, ensure that
          // the layer's position is valid. If not, mark the snapshotting as
          // failed.
          CGPoint pos = layerToRender.position;
          if (isnan(pos.x) || isnan(pos.y)) {
            snapshotSuccess = NO;
            return;
          }

          [layerToRender renderInContext:UIContext.CGContext];
        }
      }];

  if (!snapshotSuccess) {
    image = nil;
  }

  // Set the mode to UIViewTintAdjustmentModeAutomatic.
  baseView.tintAdjustmentMode = UIViewTintAdjustmentModeAutomatic;

  return image;
}

// Crops an UIImage to `frameInBaseView`.
- (UIImage*)cropImage:(UIImage*)baseImage
      frameInBaseView:(CGRect)frameInBaseView {
  if (!baseImage) {
    return nil;
  }
  DCHECK(!CGRectIsEmpty(frameInBaseView));

  // Scale `frameInBaseView` to handle an image with 2x scale.
  CGFloat scale = baseImage.scale;
  frameInBaseView.origin.x *= scale;
  frameInBaseView.origin.y *= scale;
  frameInBaseView.size.width *= scale;
  frameInBaseView.size.height *= scale;

  // Perform cropping.
  CGImageRef imageRef =
      CGImageCreateWithImageInRect(baseImage.CGImage, frameInBaseView);

  // Convert back to an UIImage.
  UIImage* image = [UIImage imageWithCGImage:imageRef
                                       scale:scale
                                 orientation:baseImage.imageOrientation];

  // Clean up a reference pointer.
  CGImageRelease(imageRef);

  return image;
}

// Returns an image of the `baseImage` overlaid with `overlays` with the given
// `frameInWindow`.
- (UIImage*)addOverlays:(NSArray<UIView*>*)overlays
              baseImage:(UIImage*)baseImage
          frameInWindow:(CGRect)frameInWindow {
  DCHECK(!CGRectIsEmpty(frameInWindow));
  if (!baseImage) {
    return nil;
  }
  // Note: If the baseImage scale differs from device scale, the baseImage size
  // may slightly differ from frameInWindow size due to rounding. Do not attempt
  // to compare the baseImage size and frameInWindow size.
  if (overlays.count == 0) {
    return baseImage;
  }
  const CGFloat kScale = [SnapshotImageScale floatImageScaleForDevice];
  DCHECK_GE(kScale, 1.0);

  UIGraphicsImageRendererFormat* format =
      [UIGraphicsImageRendererFormat preferredFormat];
  format.scale = kScale;
  format.opaque = NO;

  UIGraphicsImageRenderer* renderer =
      [[UIGraphicsImageRenderer alloc] initWithSize:frameInWindow.size
                                             format:format];

  return
      [renderer imageWithActions:^(UIGraphicsImageRendererContext* UIContext) {
        CGContextRef context = UIContext.CGContext;

        // The base image is already a cropped snapshot so it is drawn at the
        // origin of the new image.
        [baseImage drawInRect:(CGRect){.origin = CGPointZero,
                                       .size = frameInWindow.size}];

        // This shifts the origin of the context so that future drawings can be
        // in window coordinates. For example, suppose that the desired snapshot
        // area is at (0, 99) in the window coordinate space. Drawing at (0, 99)
        // will appear as (0, 0) in the resulting image.
        CGContextTranslateCTM(context, -frameInWindow.origin.x,
                              -frameInWindow.origin.y);
        [self drawOverlays:overlays context:context];
      }];
}

// Draws `overlays` onto `context` at offsets relative to the window.
- (void)drawOverlays:(NSArray<UIView*>*)overlays context:(CGContext*)context {
  for (UIView* overlay in overlays) {
    if (@available(iOS 26, *)) {
      // A translucent background starting from iOS 26 only works well with
      // drawViewHierarchyInRect.
      [overlay drawViewHierarchyInRect:overlay.bounds afterScreenUpdates:NO];
    } else {
      CGContextSaveGState(context);
      CGRect frameInWindow = [overlay.superview convertRect:overlay.frame
                                                     toView:nil];
      // This shifts the context so that drawing starts at the overlay's offset.
      CGContextTranslateCTM(context, frameInWindow.origin.x,
                            frameInWindow.origin.y);
      [[overlay layer] renderInContext:context];
      CGContextRestoreGState(context);
    }
  }
}

// Retrieves the overlays laid down on the WebState.
- (NSArray<UIView*>*)overlays {
  if (!_webState) {
    return nil;
  }
  return [_delegate
      snapshotOverlaysWithWebStateInfo:[[WebStateSnapshotInfo alloc]
                                           initWithWebState:_webState.get()]];
}

// Retrieves information needed for snapshotting.
- (std::optional<SnapshotInfo>)snapshotInfo {
  CHECK(_webState);
  SnapshotInfo snapshotInfo;
  snapshotInfo.baseView = [_delegate
      baseViewWithWebStateInfo:[[WebStateSnapshotInfo alloc]
                                   initWithWebState:_webState.get()]];
  if (!snapshotInfo.baseView) {
    return std::nullopt;
  }

  UIEdgeInsets baseViewInsets = [_delegate
      snapshotEdgeInsetsWithWebStateInfo:[[WebStateSnapshotInfo alloc]
                                             initWithWebState:_webState.get()]];
  snapshotInfo.snapshotFrameInBaseView =
      UIEdgeInsetsInsetRect(snapshotInfo.baseView.bounds, baseViewInsets);
  if (CGRectIsEmpty(snapshotInfo.snapshotFrameInBaseView)) {
    return std::nullopt;
  }

  snapshotInfo.snapshotFrameInWindow =
      [snapshotInfo.baseView convertRect:snapshotInfo.snapshotFrameInBaseView
                                  toView:nil];
  if (CGRectIsEmpty(snapshotInfo.snapshotFrameInWindow)) {
    return std::nullopt;
  }
  return snapshotInfo;
}

@end