From c04a7a8f0a95e77b62f91f197888ac73a8d77b0c Mon Sep 17 00:00:00 2001
From: Nico Hinderling <NicoHinderling@users.noreply.github.com>
Date: Mon, 26 Aug 2024 08:16:56 -0700
Subject: [PATCH] Changes

---
 .gitignore                                    |  3 +
 .../ExpandingViewController.swift             | 76 +++++++++++++++++--
 .../SnapshotPreviewsCore/View+Snapshot.swift  | 14 +++-
 3 files changed, 85 insertions(+), 8 deletions(-)

diff --git a/.gitignore b/.gitignore
index 9ec44a0..62dc33c 100644
--- a/.gitignore
+++ b/.gitignore
@@ -9,3 +9,6 @@ DerivedData/
 DemoApp/DemoApp.xcodeproj/xcuserdata/
 .xcuserstate
 PreviewsSupport/PreviewsSupport.xcframework/**/*.private.swiftinterface
+
+
+DemoApp/DemoApp.xcodeproj/project.xcworkspace/xcshareddata/swiftpm/
\ No newline at end of file
diff --git a/Sources/SnapshotPreviewsCore/ExpandingViewController.swift b/Sources/SnapshotPreviewsCore/ExpandingViewController.swift
index 62631a2..98598ff 100644
--- a/Sources/SnapshotPreviewsCore/ExpandingViewController.swift
+++ b/Sources/SnapshotPreviewsCore/ExpandingViewController.swift
@@ -41,13 +41,20 @@ extension UIView {
 
 public final class ExpandingViewController: UIHostingController<EmergeModifierView> {
 
+  // Snapshots have an overall time limit of 10s, so limit this part to 5s
+  private let HeightExpansionTimeLimitInSeconds: Double = 5
+
   private var didCall = false
   private var previousHeight: CGFloat?
 
   private var heightAnchor: NSLayoutConstraint?
   private var widthAnchor: NSLayoutConstraint?
 
-  public var expansionSettled: ((EmergeRenderingMode?, Float?, Bool?) -> Void)? {
+  private var startTime: Date?
+  private var timer: Timer?
+  private var elapsedTime: TimeInterval = 0
+
+  public var expansionSettled: ((EmergeRenderingMode?, Float?, Bool?, Error?) -> Void)? {
     didSet { didCall = false }
   }
 
@@ -90,27 +97,47 @@ public final class ExpandingViewController: UIHostingController<EmergeModifierVi
     }
   }
 
-  private func runCallback() {
+  private func runCallback(_ error: Error? = nil) {
     guard !didCall else { return }
 
     didCall = true
-    expansionSettled?(rootView.emergeRenderingMode, rootView.precision, rootView.accessibilityEnabled)
+    expansionSettled?(rootView.emergeRenderingMode, rootView.precision, rootView.accessibilityEnabled, error)
+    stopAndResetTimer()
   }
 
   public override func viewDidLayoutSubviews() {
     super.viewDidLayoutSubviews()
-    self.updateScrollViewHeight()
+
+    DispatchQueue.main.async {
+      self.updateScrollViewHeight()
+    }
   }
 
   public func updateScrollViewHeight() {
+    let supportsExpansion = rootView.supportsExpansion
+    let scrollView = view.firstScrollView
+
+    // Timeout limit
+    if timer == nil {
+      startTimer()
+    } else if elapsedTime >= HeightExpansionTimeLimitInSeconds {
+      let timeoutError = RenderingError.expandingViewTimeout(CGSize(width: UIScreen.main.bounds.size.width, height: scrollView?.visibleContentHeight ?? -1))
+      NSLog("ExpandingViewController: Expanding Scroll View timed out. Current height is \(scrollView?.visibleContentHeight ?? -1)")
+      
+      // Setting anchors back to full
+      let fittingSize = sizeThatFits(in: UIScreen.main.bounds.size)
+      heightAnchor = view.heightAnchor.constraint(greaterThanOrEqualToConstant: fittingSize.height)
+      widthAnchor = view.heightAnchor.constraint(greaterThanOrEqualToConstant: fittingSize.width)
+      runCallback(timeoutError)
+      return
+    }
+
     // If heightAnchor isn't set, this was a fixed size and we don't expand the scroll view
     guard let heightAnchor, expansionSettled != nil else {
       runCallback()
       return
     }
 
-    let supportsExpansion = rootView.supportsExpansion
-    let scrollView = view.firstScrollView
     if let scrollView, supportsExpansion {
       let diff = Int(scrollView.contentSize.height - scrollView.visibleContentHeight)
       if abs(diff) > 0 {
@@ -136,5 +163,42 @@ public final class ExpandingViewController: UIHostingController<EmergeModifierVi
     }
   }
 
+//  MARK: - Timer
+
+  func startTimer() {
+      if self.timer == nil {
+          self.startTime = Date()
+          self.timer = Timer.scheduledTimer(withTimeInterval: 0.1, repeats: true) { [weak self] _ in
+              guard let self else {
+                  return
+              }
+              if let start = self.startTime {
+                  self.elapsedTime = Date().timeIntervalSince(start)
+                  if self.elapsedTime >= HeightExpansionTimeLimitInSeconds {
+                      let scrollView = self.view.firstScrollView
+                      let timeoutError = RenderingError.expandingViewTimeout(CGSize(width: UIScreen.main.bounds.size.width, height: scrollView?.visibleContentHeight ?? -1))
+                      NSLog("ExpandingViewController: Expanding Scroll View timed out. Current height is \(scrollView?.visibleContentHeight ?? -1)")
+
+                      // Setting anchors back to full
+                      let fittingSize = self.sizeThatFits(in: UIScreen.main.bounds.size)
+                      self.heightAnchor = self.view.heightAnchor.constraint(greaterThanOrEqualToConstant: fittingSize.height)
+                      self.widthAnchor = self.view.heightAnchor.constraint(greaterThanOrEqualToConstant: fittingSize.width)
+                      self.runCallback(timeoutError)
+                  }
+              }
+          }
+          print("Timer scheduled")
+      } else {
+          print("Timer already exists")
+      }
+  }
+
+  func stopAndResetTimer() {
+      timer?.invalidate()
+      timer = nil
+      startTime = nil
+      elapsedTime = 0
+  }
+
 }
 #endif
diff --git a/Sources/SnapshotPreviewsCore/View+Snapshot.swift b/Sources/SnapshotPreviewsCore/View+Snapshot.swift
index 0af3d76..bb4892a 100644
--- a/Sources/SnapshotPreviewsCore/View+Snapshot.swift
+++ b/Sources/SnapshotPreviewsCore/View+Snapshot.swift
@@ -15,6 +15,7 @@ import SnapshotSharedModels
 public enum RenderingError: Error {
   case failedRendering(CGSize)
   case maxSize(CGSize)
+  case expandingViewTimeout(CGSize)
 }
 
 extension AccessibilityMarker: AccessibilityMark {
@@ -57,8 +58,17 @@ extension View {
     async: Bool,
     completion: @escaping (SnapshotResult) -> Void)
   {
-    controller.expansionSettled = { [weak controller, weak window] renderingMode, precision, accessibilityEnabled in
-      guard let controller, let window, let containerVC = controller.parent else { return }
+    controller.expansionSettled = { [weak controller, weak window] renderingMode, precision, accessibilityEnabled, error in
+      guard let controller, let window, let containerVC = controller.parent else {
+        return
+      }
+
+      if let error {
+        DispatchQueue.main.async {
+          completion(SnapshotResult(image: .failure(error), precision: precision, accessibilityEnabled: accessibilityEnabled, accessibilityMarkers: nil, colorScheme: _colorScheme))
+        }
+        return
+      }
 
       if async {
         DispatchQueue.main.asyncAfter(deadline: .now() + 2) {