From 799a50b39c9c2fc3c47a3e222a760113bfc0dc67 Mon Sep 17 00:00:00 2001
From: Brandon T <JustBrandonT@gmail.com>
Date: Thu, 4 Jan 2024 13:25:37 -0500
Subject: [PATCH] Fix screentime to not freeze and to set the URL for the
 back-button glitch

---
 .../Frontend/Browser/BrowserViewController.swift | 11 ++---------
 .../BrowserViewController/BVC+ScreenTime.swift   | 16 ++++++++++++----
 ...rowserViewController+TabManagerDelegate.swift | 15 +++++++++++++++
 .../BrowserViewController+ToolbarDelegate.swift  |  1 +
 4 files changed, 30 insertions(+), 13 deletions(-)

diff --git a/Sources/Brave/Frontend/Browser/BrowserViewController.swift b/Sources/Brave/Frontend/Browser/BrowserViewController.swift
index a849438f89f..19473ce92b2 100644
--- a/Sources/Brave/Frontend/Browser/BrowserViewController.swift
+++ b/Sources/Brave/Frontend/Browser/BrowserViewController.swift
@@ -1250,18 +1250,10 @@ public class BrowserViewController: UIViewController {
       make.leading.trailing.equalTo(self.view)
     }
     
-    if let screenTimeViewController = screenTimeViewController {
-      if screenTimeViewController.parent == nil {
-        addChild(screenTimeViewController)
-        screenTimeViewController.didMove(toParent: self)
-        webViewContainer.addSubview(screenTimeViewController.view)
-      }
-      
+    if let screenTimeViewController = screenTimeViewController, screenTimeViewController.parent != nil {
       screenTimeViewController.view.snp.remakeConstraints {
         $0.edges.equalTo(webViewContainer)
       }
-      
-      webViewContainer.bringSubviewToFront(screenTimeViewController.view)
     }
     
     webViewContainer.snp.remakeConstraints { make in
@@ -1948,6 +1940,7 @@ public class BrowserViewController: UIViewController {
       }
 
       updateInContentHomePanel(url as URL)
+      updateScreenTimeUrl(url)
       updatePlaylistURLBar(tab: tab, state: tab.playlistItemState, item: tab.playlistItem)
     }
   }
diff --git a/Sources/Brave/Frontend/Browser/BrowserViewController/BVC+ScreenTime.swift b/Sources/Brave/Frontend/Browser/BrowserViewController/BVC+ScreenTime.swift
index c329aa8951d..40507c4b0d9 100644
--- a/Sources/Brave/Frontend/Browser/BrowserViewController/BVC+ScreenTime.swift
+++ b/Sources/Brave/Frontend/Browser/BrowserViewController/BVC+ScreenTime.swift
@@ -12,14 +12,22 @@ extension BrowserViewController {
   /// There is also one hack required: STWebpageController breaks if you pass scheme other than http or https,
   /// it will not block anything for the rest of its lifecycle. Our internal urls have to be bridged to an empty https url.
   func updateScreenTimeUrl(_ url: URL?) {
+    guard let screenTimeViewController = screenTimeViewController else {
+      return
+    }
+    
     guard let url = url, (url.scheme == "http" || url.scheme == "https") else {
-      screenTimeViewController?.view.removeFromSuperview()
-      screenTimeViewController?.willMove(toParent: nil)
-      screenTimeViewController?.removeFromParent()
+      // This is signficantly better than removing the view controller from the screen!
+      // If we use `nil` instead, STViewController goes into a broken state PERMANENTLY until the app is restarted
+      // The URL cannot be nil, and it cannot be anything other than http(s) otherwise it will break
+      // for the duration of the app.
+      // Chromium solves this issue by not setting the URL, but instead removing it from the view entirely.
+      // But setting the URL to an empty URL works too.
+      screenTimeViewController.url = NSURL() as URL
       return
     }
     
-    screenTimeViewController?.url = url
+    screenTimeViewController.url = url
   }
   
   func recordScreenTimeUsage(for tab: Tab) {
diff --git a/Sources/Brave/Frontend/Browser/BrowserViewController/BrowserViewController+TabManagerDelegate.swift b/Sources/Brave/Frontend/Browser/BrowserViewController/BrowserViewController+TabManagerDelegate.swift
index fcf66a7f6c4..73d2657894b 100644
--- a/Sources/Brave/Frontend/Browser/BrowserViewController/BrowserViewController+TabManagerDelegate.swift
+++ b/Sources/Brave/Frontend/Browser/BrowserViewController/BrowserViewController+TabManagerDelegate.swift
@@ -64,6 +64,20 @@ extension BrowserViewController: TabManagerDelegate {
       webView.snp.remakeConstraints { make in
         make.left.right.top.bottom.equalTo(self.webViewContainer)
       }
+      
+      // Add ScreenTime above the WebView
+      if let screenTimeViewController = screenTimeViewController {
+        if screenTimeViewController.parent == nil {
+          addChild(screenTimeViewController)
+          screenTimeViewController.didMove(toParent: self)
+        }
+        
+        webViewContainer.addSubview(screenTimeViewController.view)
+        
+        screenTimeViewController.view.snp.remakeConstraints {
+          $0.edges.equalTo(webViewContainer)
+        }
+      }
 
       // This is a terrible workaround for a bad iOS 12 bug where PDF
       // content disappears any time the view controller changes (i.e.
@@ -137,6 +151,7 @@ extension BrowserViewController: TabManagerDelegate {
       topToolbar.updateReaderModeState(ReaderModeState.unavailable)
     }
 
+    updateScreenTimeUrl(tabManager.selectedTab?.url)
     updateInContentHomePanel(selected?.url as URL?)
 
     notificationsPresenter.removeNotification(with: WalletNotification.Constant.id)
diff --git a/Sources/Brave/Frontend/Browser/BrowserViewController/BrowserViewController+ToolbarDelegate.swift b/Sources/Brave/Frontend/Browser/BrowserViewController/BrowserViewController+ToolbarDelegate.swift
index 6e1c1614378..9dacdbc77bd 100644
--- a/Sources/Brave/Frontend/Browser/BrowserViewController/BrowserViewController+ToolbarDelegate.swift
+++ b/Sources/Brave/Frontend/Browser/BrowserViewController/BrowserViewController+ToolbarDelegate.swift
@@ -335,6 +335,7 @@ extension BrowserViewController: TopToolbarDelegate {
   func topToolbarDidLeaveOverlayMode(_ topToolbar: TopToolbarView) {
     hideSearchController()
     hideFavoritesController()
+    updateScreenTimeUrl(tabManager.selectedTab?.url)
     updateInContentHomePanel(tabManager.selectedTab?.url as URL?)
     updateTabsBarVisibility()
     if isUsingBottomBar {