From 8b7a59dc3cc3b507c3046e2e6b076458e45217ad Mon Sep 17 00:00:00 2001
From: Jackson Harper <jacksonh@gmail.com>
Date: Thu, 7 Dec 2023 10:36:44 +0800
Subject: [PATCH] Query for new highlights when calling findHighlight for PDFs

This fixes an issue where you could not create a note on a
new PDF highlight.
---
 apple/OmnivoreKit/Sources/App/PDFSupport/PDFViewer.swift   | 2 +-
 .../Sources/App/PDFSupport/PDFViewerViewModel.swift        | 7 +++++--
 apple/OmnivoreKit/Sources/Models/DataModels/PDFItem.swift  | 2 +-
 3 files changed, 7 insertions(+), 4 deletions(-)

diff --git a/apple/OmnivoreKit/Sources/App/PDFSupport/PDFViewer.swift b/apple/OmnivoreKit/Sources/App/PDFSupport/PDFViewer.swift
index 356a4635d..b7696be7f 100644
--- a/apple/OmnivoreKit/Sources/App/PDFSupport/PDFViewer.swift
+++ b/apple/OmnivoreKit/Sources/App/PDFSupport/PDFViewer.swift
@@ -178,7 +178,7 @@ import Utils
               if let highlight = annotations?.compactMap({ $0 as? HighlightAnnotation }).first,
                  let customHighlight = highlight.customData?["omnivoreHighlight"] as? [String: String],
                  let highlightID = customHighlight["id"]?.lowercased(),
-                 let selectedHighlight = viewModel.findHighlight(highlightID: highlightID)
+                 let selectedHighlight = viewModel.findHighlight(dataService: dataService, highlightID: highlightID)
               {
                 addNoteHighlight = selectedHighlight
                 annotation = selectedHighlight.annotation ?? ""
diff --git a/apple/OmnivoreKit/Sources/App/PDFSupport/PDFViewerViewModel.swift b/apple/OmnivoreKit/Sources/App/PDFSupport/PDFViewerViewModel.swift
index 32ca44549..c1ba16706 100644
--- a/apple/OmnivoreKit/Sources/App/PDFSupport/PDFViewerViewModel.swift
+++ b/apple/OmnivoreKit/Sources/App/PDFSupport/PDFViewerViewModel.swift
@@ -12,9 +12,11 @@ final class PDFViewerViewModel: ObservableObject {
   var snackbarMessage: String?
 
   let pdfItem: PDFItem
+  var highlights: [Highlight]
 
   init(pdfItem: PDFItem) {
     self.pdfItem = pdfItem
+    self.highlights = pdfItem.highlights
   }
 
   func snackbar(message: String) {
@@ -22,8 +24,9 @@ final class PDFViewerViewModel: ObservableObject {
     showSnackbar = true
   }
 
-  func findHighlight(highlightID: String) -> Highlight? {
-    pdfItem.highlights.first { $0.id == highlightID }
+  func findHighlight(dataService: DataService, highlightID: String) -> Highlight? {
+    let libraryItem = LibraryItem.lookup(byID: pdfItem.itemID, inContext: dataService.viewContext)
+    return libraryItem?.highlights.asArray(of: Highlight.self).first { $0.id == highlightID }
   }
 
   func loadHighlightPatches(completion onComplete: @escaping ([String]) -> Void) {
diff --git a/apple/OmnivoreKit/Sources/Models/DataModels/PDFItem.swift b/apple/OmnivoreKit/Sources/Models/DataModels/PDFItem.swift
index 3dcc5c39e..4b5601d9e 100644
--- a/apple/OmnivoreKit/Sources/Models/DataModels/PDFItem.swift
+++ b/apple/OmnivoreKit/Sources/Models/DataModels/PDFItem.swift
@@ -15,7 +15,7 @@ public struct PDFItem {
   public let isArchived: Bool
   public let isRead: Bool
   public let downloadURL: String
-  public let highlights: [Highlight]
+  public var highlights: [Highlight]
 
   public static func make(item: LibraryItem) -> PDFItem? {
     guard item.isPDF else { return nil }