// Copyright 2024 The Chromium Authors
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

import UIKit

// A class that generates snapshot images for the WebState associated with this class.
@MainActor
@objcMembers public class SnapshotGenerator: NSObject {
  // A wrapper class for the associated WebState.
  private let webStateInfo: WebStateSnapshotInfo
  // The SnapshotGenerator delegate to obtain the information about UIView.
  weak var delegate: SnapshotGeneratorDelegate?

  // Designated initializer.
  init(webStateInfo: WebStateSnapshotInfo) {
    self.webStateInfo = webStateInfo
    self.delegate = nil
  }

  // Generates a new snapshot and runs a callback with the new snapshot image.
  func generateSnapshot(completion: ((UIImage?) -> Void)?) {
    generateSnapshot(includeOverlays: true, completion: completion)
  }

  // Generates a new snapshot and runs a callback with the new snapshot image.
  // The generated image includes overlays if `includeOverlays` is true.
  func generateSnapshot(includeOverlays: Bool, completion: ((UIImage?) -> Void)?) {
    guard delegate != nil else {
      completion?(nil)
      return
    }
    guard let lastCommittedUrl = webStateInfo.lastCommittedURL(),
      let newTabPageUrl = URL.init(string: "chrome://newtab/")
    else {
      completion?(nil)
      return
    }
    // 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 (generateWKWebViewSnapshot() vs generateUIViewSnapshotWithOverlays()).
    let isNTP = lastCommittedUrl.nsurl == newTabPageUrl
    if !isNTP && webStateInfo.canTakeSnapshot() {
      // Take the snapshot using the optimized WKWebView snapshotting API for pages loaded in the
      // web view when the WebState snapshot API is available.
      generateWKWebViewSnapshot(includeOverlays: includeOverlays, completion: completion)
      return
    }
    // Use the UIKit-based snapshot API as a fallback when the WKWebView API is unavailable.
    let snapshot = includeOverlays ? generateUIViewSnapshotWithOverlays() : generateUIViewSnapshot()
    if completion != nil {
      completion?(snapshot)
    }
  }

  // Generates and returns a new snapshot image with UIKit-based snapshot API.
  func generateUIViewSnapshot() -> UIImage? {
    if !canTakeSnapshot() {
      return nil
    }

    guard let delegate = delegate else {
      return nil
    }

    delegate.willUpdateSnapshot(webStateInfo: webStateInfo)

    guard let baseView = delegate.baseView(webStateInfo: webStateInfo) else {
      return nil
    }
    let baseViewInsets = delegate.snapshotEdgeInsets(webStateInfo: webStateInfo)
    let frameInBaseView = baseView.bounds.inset(by: baseViewInsets)

    let baseImage = convertBaseView(baseView)
    return cropImage(baseImage: baseImage, frameInBaseView: frameInBaseView)
  }

  // Generates and returns a new snapshot image with UIKit-based snapshot API. The generated image
  // includes overlays (e.g., infobars, the download manager, and sad tab view).
  func generateUIViewSnapshotWithOverlays() -> UIImage? {
    return addOverlays(baseImage: generateUIViewSnapshot())
  }

  // 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.
  private func generateWKWebViewSnapshot(includeOverlays: Bool, completion: ((UIImage?) -> Void)?) {
    if !canTakeSnapshot() {
      completion?(nil)
      return
    }
    guard let delegate = delegate else {
      completion?(nil)
      return
    }
    delegate.willUpdateSnapshot(webStateInfo: webStateInfo)

    guard let baseView = delegate.baseView(webStateInfo: webStateInfo) else {
      completion?(nil)
      return
    }
    let baseViewInsets = delegate.snapshotEdgeInsets(webStateInfo: webStateInfo)
    let frameInBaseView = baseView.bounds.inset(by: baseViewInsets)

    let wrappedCompletion = { [weak self] (image: UIImage?) in
      let snapshot = self?.adjustWKWebViewSnapshotIfNecessary(
        image: image, includeOverlays: includeOverlays)
      completion?(snapshot)
    }
    webStateInfo.takeSnapshot(frameInBaseView, callback: wrappedCompletion)
  }

  // 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 true.
  private func adjustWKWebViewSnapshotIfNecessary(image: UIImage?, includeOverlays: Bool)
    -> UIImage?
  {
    guard let image = image else {
      return nil
    }
    guard let delegate = delegate else {
      return nil
    }
    guard let baseView = delegate.baseView(webStateInfo: webStateInfo)
    else {
      return nil
    }
    let baseViewInsets = delegate.snapshotEdgeInsets(webStateInfo: webStateInfo)
    let frameInBaseView = baseView.bounds.inset(by: baseViewInsets)

    // 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 {
      guard let backgroundImage = self.generateUIViewSnapshot() else {
        return nil
      }

      let format = UIGraphicsImageRendererFormat.preferred()
      format.scale = SnapshotImageScale.floatForDevice()
      format.opaque = true

      let renderer = UIGraphicsImageRenderer(size: frameInBaseView.size, format: format)
      let combinedImage = autoreleasepool {
        return renderer.image { (context) in
          backgroundImage.draw(in: CGRect(origin: CGPoint.zero, size: backgroundImage.size))
          image.draw(in: CGRect(origin: CGPoint.zero, size: image.size))
        }
      }

      return includeOverlays ? self.addOverlays(baseImage: combinedImage) : combinedImage
    }

    return includeOverlays ? self.addOverlays(baseImage: image) : image
  }

  // Converts an UIView to an UIImage. The size of generated UIImage is the same as `baseView`.
  private func convertBaseView(_ baseView: UIView) -> UIImage? {
    // Disable the automatic view dimming UIKit performs if a view is presented modally over
    // `baseView`.
    baseView.tintAdjustmentMode = .normal

    let format = UIGraphicsImageRendererFormat.preferred()
    format.scale = SnapshotImageScale.floatForDevice()
    format.opaque = false

    let renderer = UIGraphicsImageRenderer(bounds: baseView.bounds, format: format)

    var snapshotSuccess = true
    let image = autoreleasepool {
      return renderer.image { context in
        if #available(iOS 26, *) {
          // A translucent background starting from iOS 26 only works well with drawHierarchy.
          snapshotSuccess = baseView.drawHierarchy(in: baseView.bounds, afterScreenUpdates: false)
        } else {
          // Take animations into account by rendering the presentation layer.
          // Fallback to the rendering the layer if not possible.
          let layerToRender = baseView.layer.presentation() ?? baseView.layer

          // To mitigate against crashes like crbug.com/1429512, ensure that the
          // layer's position is valid. Otherwise mark the snapshotting as failed.
          let position = layerToRender.position
          let validPosition = !position.x.isNaN && !position.y.isNaN
          guard validPosition else {
            snapshotSuccess = false
            return
          }

          layerToRender.render(in: context.cgContext)
        }
      }
    }

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

    return snapshotSuccess ? image : nil
  }

  // Crops an UIImage to `frameInBaseView`.
  private func cropImage(baseImage: UIImage?, frameInBaseView: CGRect) -> UIImage? {
    guard let baseImage = baseImage else {
      return nil
    }
    guard let cgImage = baseImage.cgImage else {
      return nil
    }
    let scale = baseImage.scale

    var frame = frameInBaseView
    frame.origin.x *= scale
    frame.origin.y *= scale
    frame.size.width *= scale
    frame.size.height *= scale
    guard let croppedCGImage = cgImage.cropping(to: frame) else {
      return nil
    }

    return UIImage(
      cgImage: croppedCGImage, scale: baseImage.imageRendererFormat.scale,
      orientation: baseImage.imageOrientation)
  }

  // Returns false if WebState or the view is not ready for snapshot.
  private func canTakeSnapshot() -> Bool {
    // This allows for easier unit testing of classes that use SnapshotGenerator.
    guard let delegate = delegate else {
      return false
    }

    // Do not generate a snapshot if web usage is disabled (as the WebState's view is blank in that
    // case).
    if !webStateInfo.isWebUsageEnabled() {
      return false
    }

    return delegate.canTakeSnapshot(webStateInfo: webStateInfo)
  }

  // Returns an image of the `baseImage` overlaid with overlays.
  private func addOverlays(baseImage: UIImage?) -> UIImage? {
    guard let delegate = delegate else {
      return nil
    }
    guard let baseImage = baseImage else {
      return nil
    }
    guard let baseView = delegate.baseView(webStateInfo: webStateInfo) else {
      return nil
    }
    let baseViewInsets = delegate.snapshotEdgeInsets(webStateInfo: webStateInfo)
    let frameInBaseView = baseView.bounds.inset(by: baseViewInsets)
    let snapshotFrameInWindow = baseView.convert(frameInBaseView, to: nil)

    let overlays = delegate.snapshotOverlays(webStateInfo: webStateInfo)
    // Note: If the baseImage scale differs from device scale, the baseImage size may slightly
    // differ from `snapshotFrameInWindow` size due to rounding. Do not attempt to compare the
    // `baseImage` size and `snapshotFrameInWindow` size.
    if overlays.count == 0 {
      return baseImage
    }

    let format = UIGraphicsImageRendererFormat.preferred()
    format.scale = SnapshotImageScale.floatForDevice()
    format.opaque = false

    let renderer = UIGraphicsImageRenderer(size: snapshotFrameInWindow.size, format: format)

    return autoreleasepool {
      return renderer.image { (context) in
        let cgContextRef = context.cgContext

        // The base image is already a cropped snapshot so it is drawn at the origin of the new image.
        baseImage.draw(in: CGRect(origin: CGPoint.zero, size: snapshotFrameInWindow.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.
        cgContextRef.translateBy(
          x: -snapshotFrameInWindow.origin.x, y: -snapshotFrameInWindow.origin.y)

        for overlay in overlays {
          if #available(iOS 26, *) {
            // A translucent background starting from iOS 26 only works well with drawHierarchy.
            overlay.drawHierarchy(in: overlay.bounds, afterScreenUpdates: false)
          } else {
            cgContextRef.saveGState()
            let superview = overlay.superview ?? baseView
            // The following is only correct if `superview` is indeed `overlay.superview`, since `frame`
            // is defined as "the view’s location and size in its superview’s coordinate system".
            // However it appears that some overlays do not have any superviews.
            let frameInWindow = superview.convert(overlay.frame, to: nil)
            overlay.layer.render(in: cgContextRef)
            // This shifts the context so that drawing starts at the overlay's offset.
            cgContextRef.translateBy(x: frameInWindow.origin.x, y: frameInWindow.origin.y)
            cgContextRef.restoreGState()
          }
        }
      }
    }
  }
}