Tusker/Packages/ComposeUI/Sources/ComposeUI/CoreData/DraftsPersistentContainer.swift

//
//  DraftsPersistentContainer.swift
//  ComposeUI
//
//  Created by Shadowfacts on 4/22/23.
//

import Foundation
import CoreData
import OSLog
import Pachyderm

private let logger = Logger(subsystem: Bundle.main.bundleIdentifier!, category: "DraftsPersistentContainer")

public class DraftsPersistentContainer: NSPersistentContainer {
    
    public static let shared = DraftsPersistentContainer()
    
    private static let managedObjectModel: NSManagedObjectModel = {
        let url = Bundle.module.url(forResource: "Drafts", withExtension: "momd")!
        return NSManagedObjectModel(contentsOf: url)!
    }()
    
    private var lastHistoryToken: NSPersistentHistoryToken!
    
    init() {
        super.init(name: "Drafts", managedObjectModel: DraftsPersistentContainer.managedObjectModel)
        
        let containerURL = FileManager.default.containerURL(forSecurityApplicationGroupIdentifier: "group.space.vaccor.Tusker")!
        let documentsURL = containerURL.appendingPathComponent("Documents")
        let storeDesc = NSPersistentStoreDescription(url: documentsURL.appendingPathComponent("drafts").appendingPathExtension("sqlite"))
        storeDesc.type = NSSQLiteStoreType
        storeDesc.setOption(true as NSNumber, forKey: NSPersistentHistoryTrackingKey)
        storeDesc.setOption(true as NSNumber, forKey: NSPersistentStoreRemoteChangeNotificationPostOptionKey)
        
        persistentStoreDescriptions = [
            storeDesc
        ]
        
        loadPersistentStores { _, error in
            if let error {
                fatalError("Loading persistent store: \(error)")
            }
        }
        
        viewContext.automaticallyMergesChangesFromParent = true
        viewContext.mergePolicy = NSMergePolicy.mergeByPropertyObjectTrump
        
        lastHistoryToken = persistentStoreCoordinator.currentPersistentHistoryToken(fromStores: nil)
        
        NotificationCenter.default.addObserver(self, selector: #selector(remoteChanges(_:)), name: .NSPersistentStoreRemoteChange, object: persistentStoreCoordinator)
    }
    
    public func save() {
        guard viewContext.hasChanges else {
            return
        }
        do {
            try viewContext.save()
        } catch {
            logger.error("Failed to save: \(String(describing: error))")
        }
    }
    
    public func migrate(from url: URL, completion: @escaping (Result<(), any Error>) -> Void) {
        performBackgroundTask { context in
            let result = DraftsMigrator.migrate(from: url, to: context)
            completion(result)
            try! context.save()
        }
    }
    
    public func getDraft(id: UUID) -> Draft? {
        let req = Draft.fetchRequest(id: id)
        return try? viewContext.fetch(req).first
    }
    
    public func createDraft(
        accountID: String,
        text: String,
        contentWarning: String,
        inReplyToID: String?,
        visibility: Visibility,
        localOnly: Bool
    ) -> Draft {
        let draft = Draft(context: viewContext)
        draft.accountID = accountID
        draft.text = text
        draft.initialText = text
        draft.contentWarning = contentWarning
        draft.contentWarningEnabled = !contentWarning.isEmpty
        draft.inReplyToID = inReplyToID
        draft.visibility = visibility
        draft.localOnly = localOnly
        save()
        return draft
    }
    
    public func createEditDraft(
        accountID: String,
        source: StatusSource,
        inReplyToID: String?,
        visibility: Visibility,
        localOnly: Bool,
        attachments: [Attachment],
        poll: Pachyderm.Poll?
    ) -> Draft {
        let draft = Draft(context: viewContext)
        draft.accountID = accountID
        draft.editedStatusID = source.id
        draft.text = source.text
        draft.initialText = source.text
        draft.contentWarning = source.spoilerText
        draft.contentWarningEnabled = !source.spoilerText.isEmpty
        draft.inReplyToID = inReplyToID
        draft.visibility = visibility
        draft.localOnly = localOnly
        for attachment in attachments {
            createEditDraftAttachment(attachment, in: draft)
        }
        if let existingPoll = poll {
            let poll = Poll(context: viewContext)
            poll.draft = draft
            draft.poll = poll
            if let expiresAt = existingPoll.expiresAt,
               !existingPoll.effectiveExpired {
                poll.duration = PollController.Duration.allCases.max(by: {
                    (expiresAt.timeIntervalSinceNow - $0.timeInterval) < (expiresAt.timeIntervalSinceNow - $1.timeInterval)
                })!.timeInterval
            } else {
                poll.duration = PollController.Duration.oneDay.timeInterval
            }
            poll.multiple = existingPoll.multiple
            // rmeove default empty options
            for opt in poll.pollOptions {
                viewContext.delete(opt)
            }
            for existingOpt in existingPoll.options {
                let opt = PollOption(context: viewContext)
                opt.poll = poll
                poll.options.add(opt)
                opt.text = existingOpt.title
            }
        }
        save()
        return draft
    }
    
    private func createEditDraftAttachment(_ attachment: Attachment, in draft: Draft) {
        let draftAttachment = DraftAttachment(context: viewContext)
        draftAttachment.id = UUID()
        draftAttachment.attachmentDescription = attachment.description ?? ""
        draftAttachment.editedAttachmentID = attachment.id
        draftAttachment.editedAttachmentKind = attachment.kind
        draftAttachment.editedAttachmentURL = attachment.url
        draftAttachment.draft = draft
        draft.attachments.add(draftAttachment)
    }
    
    public func removeOrphanedAttachments(completion: @escaping () -> Void) {
        guard let files = try? FileManager.default.contentsOfDirectory(at: DraftAttachment.attachmentsDirectory, includingPropertiesForKeys: nil),
              !files.isEmpty else {
            return
        }
        performBackgroundTask { context in
            let allAttachmentsReq = DraftAttachment.fetchRequest()
            allAttachmentsReq.predicate = NSPredicate(format: "fileURL != nil")
            guard let allAttachments = try? context.fetch(allAttachmentsReq) else {
                return
            }
            let orphaned = Set(files).subtracting(allAttachments.lazy.compactMap(\.fileURL))
            for url in orphaned {
                do {
                    try FileManager.default.removeItem(at: url)
                } catch {
                    logger.error("Failed to remove orphaned attachment: \(String(describing: error), privacy: .public)")
                }
            }
            completion()
        }
    }
    
    @objc private func remoteChanges(_ notification: Foundation.Notification) {
        guard let newHistoryToken = notification.userInfo?[NSPersistentHistoryTokenKey] as? NSPersistentHistoryToken else {
            return
        }
        
        // todo: should this be on a background context?
        let context = viewContext
        context.perform {
            let predicate = NSPredicate(format: "(%@ < token) AND (token <= %@)", self.lastHistoryToken, newHistoryToken)
            
            let historyRequest = NSPersistentHistoryTransaction.fetchRequest!
            historyRequest.predicate = predicate
            let request = NSPersistentHistoryChangeRequest.fetchHistory(withFetch: historyRequest)
            if let result = try? context.execute(request) as? NSPersistentHistoryResult,
               let transactions = result.result as? [NSPersistentHistoryTransaction] {
                for transaction in transactions {
                    guard let userInfo = transaction.objectIDNotification().userInfo else {
                        continue
                    }
                    NSManagedObjectContext.mergeChanges(fromRemoteContextSave: userInfo, into: [context])
                }
            }
            
            self.lastHistoryToken = newHistoryToken
        }
    }
    
}