Add SwiftData-backed image support

TakeNote/Library/NoteImageManager.swift

@@ -0,0 +1,130 @@
+//
+//  NoteImageManager.swift
+//  TakeNote
+//
+//  Created by Adam Drew on 9/9/25.
+//
+
+import SwiftData
+import SwiftUI
+import UniformTypeIdentifiers
+
+@MainActor
+@Observable
+class NoteImageManager {
+    var modelContext: ModelContext
+
+    init(modelContext: ModelContext) {
+        self.modelContext = modelContext
+    }
+
+    func ingestImage(from url: URL, note: Note) -> NoteImage? {
+        guard url.isFileURL else { return nil }
+        guard let type = UTType(filenameExtension: url.pathExtension),
+            type.conforms(to: .image)
+        else { return nil }
+        guard let data = try? Data(contentsOf: url) else { return nil }
+
+        let mimeType = type.preferredMIMEType ?? "application/octet-stream"
+        let image = NoteImage(data: data, mimeType: mimeType, referenceCount: 1)
+        modelContext.insert(image)
+        let link = NoteImageLink(sourceNote: note, image: image)
+        modelContext.insert(link)
+        return image
+    }
+
+    func updateImageLinks(for note: Note) {
+        let existingLinks = getLinksForSourceNote(note)
+        let existingUUIDs = Set(
+            existingLinks.compactMap { $0.image?.uuid }
+        )
+        let currentUUIDs = Set(extractImageUUIDs(from: note.content))
+
+        let linksToRemove = existingLinks.filter { link in
+            guard let uuid = link.image?.uuid else { return true }
+            return !currentUUIDs.contains(uuid)
+        }
+
+        for link in linksToRemove {
+            if let image = link.image {
+                image.referenceCount = max(0, image.referenceCount - 1)
+                if image.referenceCount == 0 {
+                    modelContext.delete(image)
+                }
+            }
+            modelContext.delete(link)
+        }
+
+        let uuidsToAdd = currentUUIDs.subtracting(existingUUIDs)
+        if !uuidsToAdd.isEmpty {
+            let images = getImagesForUUIDs(uuidsToAdd)
+            let imageByUUID = makeUUIDImageMap(images)
+            for uuid in uuidsToAdd {
+                guard let image = imageByUUID[uuid] else { continue }
+                let link = NoteImageLink(sourceNote: note, image: image)
+                image.referenceCount += 1
+                modelContext.insert(link)
+            }
+        }
+
+        try? modelContext.save()
+    }
+
+    func removeImageLinks(for note: Note) {
+        let existingLinks = getLinksForSourceNote(note)
+        for link in existingLinks {
+            if let image = link.image {
+                image.referenceCount = max(0, image.referenceCount - 1)
+                if image.referenceCount == 0 {
+                    modelContext.delete(image)
+                }
+            }
+            modelContext.delete(link)
+        }
+        try? modelContext.save()
+    }
+
+    private func extractImageUUIDs(from text: String) -> [UUID] {
+        let pattern =
+            #/(?i)takenote:\/\/image\/([0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12})/#
+        var seen = Set<UUID>()
+        var result: [UUID] = []
+
+        for match in text.matches(of: pattern) {
+            let uuidSub = match.1
+            if let uuid = UUID(uuidString: String(uuidSub)),
+                seen.insert(uuid).inserted
+            {
+                result.append(uuid)
+            }
+        }
+        return result
+    }
+
+    private func getLinksForSourceNote(_ note: Note) -> [NoteImageLink] {
+        let uuid = note.uuid
+        return
+            (try? modelContext.fetch(
+                FetchDescriptor<NoteImageLink>(
+                    predicate: #Predicate { $0.sourceNote?.uuid == uuid }
+                )
+            )) ?? []
+    }
+
+    private func getImagesForUUIDs(_ uuids: Set<UUID>) -> [NoteImage] {
+        return
+            (try? modelContext.fetch(
+                FetchDescriptor<NoteImage>(
+                    predicate: #Predicate { uuids.contains($0.uuid) }
+                )
+            )) ?? []
+    }
+
+    private func makeUUIDImageMap(_ images: [NoteImage]) -> [UUID: NoteImage] {
+        var imageByUUID: [UUID: NoteImage] = [:]
+        for image in images {
+            imageByUUID[image.uuid] = image
+        }
+        return imageByUUID
+    }
+}

TakeNote/Library/TakeNoteImageURLProtocol.swift

@@ -0,0 +1,85 @@
+//
+//  TakeNoteImageURLProtocol.swift
+//  TakeNote
+//
+//  Created by Adam Drew on 9/9/25.
+//
+
+import Foundation
+import SwiftData
+
+enum TakeNoteImageURLProtocolRegistrar {
+    private static var isRegistered = false
+
+    static func registerIfNeeded(container: ModelContainer) {
+        guard !isRegistered else { return }
+        TakeNoteImageURLProtocol.modelContainer = container
+        URLProtocol.registerClass(TakeNoteImageURLProtocol.self)
+        isRegistered = true
+    }
+}
+
+final class TakeNoteImageURLProtocol: URLProtocol {
+    static var modelContainer: ModelContainer?
+
+    override class func canInit(with request: URLRequest) -> Bool {
+        guard let url = request.url else { return false }
+        return url.scheme == "takenote" && url.host == "image"
+    }
+
+    override class func canonicalRequest(for request: URLRequest) -> URLRequest {
+        return request
+    }
+
+    override func startLoading() {
+        guard let url = request.url else { return }
+        guard let uuid = UUID(uuidString: url.lastPathComponent) else {
+            client?.urlProtocol(
+                self,
+                didFailWithError: URLError(.badURL)
+            )
+            return
+        }
+        guard let container = Self.modelContainer else {
+            client?.urlProtocol(
+                self,
+                didFailWithError: URLError(.resourceUnavailable)
+            )
+            return
+        }
+
+        Task { [weak self] @MainActor in
+            guard let self else { return }
+            let context = ModelContext(container)
+            let image = (try? context.fetch(
+                FetchDescriptor<NoteImage>(
+                    predicate: #Predicate { $0.uuid == uuid }
+                )
+            ))?.first
+
+            guard let image else {
+                self.client?.urlProtocol(
+                    self,
+                    didFailWithError: URLError(.fileDoesNotExist)
+                )
+                return
+            }
+
+            let response = URLResponse(
+                url: url,
+                mimeType: image.mimeType,
+                expectedContentLength: image.data.count,
+                textEncodingName: nil
+            )
+            self.client?.urlProtocol(
+                self,
+                didReceive: response,
+                cacheStoragePolicy: .notAllowed
+            )
+            self.client?.urlProtocol(self, didLoad: image.data)
+            self.client?.urlProtocolDidFinishLoading(self)
+        }
+    }
+
+    override func stopLoading() {}
+}

TakeNote/Models/Note.swift

@@ -75,6 +75,7 @@ class Note: Identifiable {
     // (we'll specify inverses on NoteLink to avoid macro circularity).
     @Relationship var outgoingLinks: [NoteLink]? = []
     @Relationship var incomingLinks: [NoteLink]? = []
+    @Relationship var imageLinks: [NoteImageLink]? = []
 
     init(folder: NoteContainer) {
         self.title = self.defaultTitle

TakeNote/Models/NoteImage.swift

@@ -0,0 +1,38 @@
+//
+//  NoteImage.swift
+//  TakeNote
+//
+//  Created by Adam Drew on 9/9/25.
+//
+
+import Foundation
+import SwiftData
+import SwiftUI
+
+// Hey!
+// Hey you!
+// If you change model schema remember to bump ckBootstrapVersionCurrent
+// in TakeNoteApp.swift
+//
+// And don't forget to promote to prod!!!
+
+@Model
+class NoteImage {
+    private(set) var uuid: UUID = UUID()
+    @Attribute(.externalStorage) var data: Data
+    var mimeType: String
+    var referenceCount: Int
+    var createdDate: Date = Date()
+    @Relationship var noteLinks: [NoteImageLink]? = []
+
+    init(data: Data, mimeType: String, referenceCount: Int = 0) {
+        self.data = data
+        self.mimeType = mimeType
+        self.referenceCount = referenceCount
+        self.uuid = UUID()
+    }
+
+    func getURL() -> String {
+        return "takenote://image/\(uuid.uuidString)"
+    }
+}

TakeNote/Models/NoteImageLink.swift

@@ -0,0 +1,29 @@
+//
+//  NoteImageLink.swift
+//  TakeNote
+//
+//  Created by Adam Drew on 9/9/25.
+//
+
+import SwiftData
+import SwiftUI
+
+// Hey!
+// Hey you!
+// If you change model schema remember to bump ckBootstrapVersionCurrent
+// in TakeNoteApp.swift
+//
+// And don't forget to promote to prod!!!
+
+@Model
+class NoteImageLink {
+    @Relationship(inverse: \Note.imageLinks)
+    var sourceNote: Note?
+    @Relationship(inverse: \NoteImage.noteLinks)
+    var image: NoteImage?
+
+    init(sourceNote: Note, image: NoteImage) {
+        self.sourceNote = sourceNote
+        self.image = image
+    }
+}

TakeNote/TakeNoteApp.swift

@@ -11,7 +11,7 @@ private let onboardingVersionKey = "onboarding.version.seen"
 
 #if DEBUG
     // Bump this to get the schema to update, for example if there have been model changes
-    private let ckBootstrapVersionCurrent = 8
+    private let ckBootstrapVersionCurrent = 9
     private let ckBootstrapVersionKey = "takenote.ck.bootstrap.version"
 #endif
 
@@ -73,7 +73,13 @@ struct TakeNoteApp: App {
                     "CKBootstrap-\(UUID().uuidString).sqlite"
                 )
             AppBootstrapper.bootstrapDevSchemaIfNeeded(
-                modelTypes: [Note.self, NoteContainer.self, NoteLink.self],
+                modelTypes: [
+                    Note.self,
+                    NoteContainer.self,
+                    NoteLink.self,
+                    NoteImage.self,
+                    NoteImageLink.self
+                ],
                 storeURL: tempBootstrapURL,
                 containerID: "iCloud.com.adamdrew.takenote",
                 userDefaultsKey: ckBootstrapVersionKey,
@@ -88,12 +94,18 @@ struct TakeNoteApp: App {
                 for: Note.self,
                 NoteContainer.self,
                 NoteLink.self,
+                NoteImage.self,
+                NoteImageLink.self,
                 configurations: config
             )
         } catch {
             fatalError("Failed to initialize ModelContainer: \(error)")
         }
 
+        TakeNoteImageURLProtocolRegistrar.registerIfNeeded(
+            container: container
+        )
+
         // Capture values in locals to avoid capturing `self` in escaping closures
         let modelContainer = container
         let viewModel = takeNoteVM

TakeNote/TakeNoteVM.swift

@@ -284,10 +284,12 @@ class TakeNoteVM {
             errorAlertIsVisible = true
             return
         }
+        let imageManager = NoteImageManager(modelContext: modelContext)
         for note in trash.notes {
             if openNote == note {
                 openNote = nil
             }
+            imageManager.removeImageLinks(for: note)
             modelContext.delete(note)
         }
         do {

TakeNote/Views/MainWindow/MainWindow.swift

@@ -22,6 +22,8 @@ struct MainWindow: View {
     @Query() var notes: [Note]
     @Query() var containers: [NoteContainer]
     @Query() var noteLinks: [NoteLink]
+    @Query() var noteImageLinks: [NoteImageLink]
+    @Query() var noteImages: [NoteImage]
 
     @State var notesInBufferMessagePresented: Bool = false
     @State var showDeleteEverythingAlert: Bool = false
@@ -71,6 +73,12 @@ struct MainWindow: View {
         for noteLink in noteLinks {
             modelContext.delete(noteLink)
         }
+        for imageLink in noteImageLinks {
+            modelContext.delete(imageLink)
+        }
+        for image in noteImages {
+            modelContext.delete(image)
+        }
         try? modelContext.save()
     }