From 851f8f2a991c11b6ecdf799714c603bbe3334934 Mon Sep 17 00:00:00 2001
From: mark <saybyebyet@gmail.com>
Date: Tue, 30 Mar 2021 21:24:09 +0800
Subject: [PATCH] Reset with first album when there're no album after user
 granted some of them

---
 .../ImagePickerController+Albums.swift        |  1 +
 .../Controller/ImagePickerController.swift    | 33 ++++++++++++++++++-
 2 files changed, 33 insertions(+), 1 deletion(-)

diff --git a/Sources/Controller/ImagePickerController+Albums.swift b/Sources/Controller/ImagePickerController+Albums.swift
index a5020cb1..5a75f8db 100644
--- a/Sources/Controller/ImagePickerController+Albums.swift
+++ b/Sources/Controller/ImagePickerController+Albums.swift
@@ -34,6 +34,7 @@ extension ImagePickerController: AlbumsViewControllerDelegate {
     }
 
     func select(album: PHAssetCollection) {
+        self.currentAlbum = album
         assetsViewController.showAssets(in: album)
         albumButton.setTitle((album.localizedTitle ?? "") + " ", for: .normal)
         albumButton.sizeToFit()
diff --git a/Sources/Controller/ImagePickerController.swift b/Sources/Controller/ImagePickerController.swift
index 85b086de..98fdc6af 100644
--- a/Sources/Controller/ImagePickerController.swift
+++ b/Sources/Controller/ImagePickerController.swift
@@ -56,7 +56,7 @@ import Photos
     let dropdownTransitionDelegate = DropdownTransitionDelegate()
     let zoomTransitionDelegate = ZoomTransitionDelegate()
 
-    lazy var albums: [PHAssetCollection] = {
+    private func loadAlbums() -> [PHAssetCollection] {
         // We don't want collections without assets.
         // I would like to do that with PHFetchOptions: fetchOptions.predicate = NSPredicate(format: "estimatedAssetCount > 0")
         // But that doesn't work...
@@ -74,6 +74,11 @@ import Photos
             let assetsFetchResult = PHAsset.fetchAssets(in: $0, options: fetchOptions)
             return assetsFetchResult.count > 0
         }
+    }
+    var currentAlbum: PHAssetCollection?
+
+    lazy var albums: [PHAssetCollection] = {
+        self.loadAlbums()
     }()
 
     public init(selectedAssets: [PHAsset] = []) {
@@ -86,6 +91,16 @@ import Photos
         fatalError("init(coder:) has not been implemented")
     }
 
+    open override func viewWillAppear(_ animated: Bool) {
+        super.viewWillAppear(animated)
+        PHPhotoLibrary.shared().register(self)
+    }
+
+    open override func viewWillDisappear(_ animated: Bool) {
+        super.viewWillDisappear(animated)
+        PHPhotoLibrary.shared().unregisterChangeObserver(self)
+    }
+
     public override func viewDidLoad() {
         super.viewDidLoad()
         
@@ -161,3 +176,19 @@ import Photos
         albumButton.isHidden = albums.count < 2
     }
 }
+
+extension ImagePickerController: PHPhotoLibraryChangeObserver {
+    public func photoLibraryDidChange(_ changeInstance: PHChange) {
+        if let currentAlbum = currentAlbum,
+           changeInstance.changeDetails(for: currentAlbum) != nil {
+            return
+        }
+        DispatchQueue.main.async {
+            self.albums = self.loadAlbums()
+            self.updateAlbumButton()
+            if let firstAlbum = self.albums.first {
+                self.select(album: firstAlbum)
+            }
+        }
+    }
+}