InvestmentTrackerApp/PortfolioJournal/Models/CoreDataStack.swift

import CoreData
import CloudKit
import Combine
import WidgetKit

class CoreDataStack: ObservableObject {
    static let shared = CoreDataStack()

    static let appGroupIdentifier = "group.com.alexandrevazquez.portfoliojournal"
    static let cloudKitContainerIdentifier = "iCloud.com.alexandrevazquez.portfoliojournal"

    private static var cloudKitEnabled: Bool {
        UserDefaults.standard.bool(forKey: "cloudSyncEnabled")
    }

    private static var appGroupEnabled: Bool {
        true
    }

    private static func migrateStoreIfNeeded(from sourceURL: URL, to destinationURL: URL) {
        let fileManager = FileManager.default
        guard fileManager.fileExists(atPath: sourceURL.path) else {
            return
        }

        let sourceHasData = storeHasData(at: sourceURL)
        let destinationHasData = storeHasData(at: destinationURL)
        guard sourceHasData, !destinationHasData else { return }

        removeStoreFilesIfNeeded(at: destinationURL)

        let relatedSuffixes = ["", "-wal", "-shm"]
        for suffix in relatedSuffixes {
            let source = URL(fileURLWithPath: sourceURL.path + suffix)
            let destination = URL(fileURLWithPath: destinationURL.path + suffix)
            guard fileManager.fileExists(atPath: source.path),
                  !fileManager.fileExists(atPath: destination.path) else {
                continue
            }

            do {
                try fileManager.copyItem(at: source, to: destination)
            } catch {
                print("Core Data store migration failed for \(source.lastPathComponent): \(error)")
            }
        }
    }

    private static func removeStoreFilesIfNeeded(at url: URL) {
        let fileManager = FileManager.default
        let relatedSuffixes = ["", "-wal", "-shm"]
        for suffix in relatedSuffixes {
            let fileURL = URL(fileURLWithPath: url.path + suffix)
            guard fileManager.fileExists(atPath: fileURL.path) else { continue }
            do {
                try fileManager.removeItem(at: fileURL)
            } catch {
                print("Failed to remove existing store file \(fileURL.lastPathComponent): \(error)")
            }
        }
    }

    private static func storeHasData(at url: URL) -> Bool {
        let fileManager = FileManager.default
        guard fileManager.fileExists(atPath: url.path),
              let model = NSManagedObjectModel.mergedModel(from: [Bundle.main]) else {
            return false
        }

        let coordinator = NSPersistentStoreCoordinator(managedObjectModel: model)
        do {
            try coordinator.addPersistentStore(
                ofType: NSSQLiteStoreType,
                configurationName: nil,
                at: url,
                options: [NSReadOnlyPersistentStoreOption: true]
            )
        } catch {
            print("Failed to open store for data check: \(error)")
            return false
        }

        let context = NSManagedObjectContext(concurrencyType: .privateQueueConcurrencyType)
        context.persistentStoreCoordinator = coordinator

        let snapshotRequest = NSFetchRequest<NSManagedObject>(entityName: "Snapshot")
        snapshotRequest.resultType = .countResultType
        let snapshotCount = (try? context.count(for: snapshotRequest)) ?? 0
        if snapshotCount > 0 {
            return true
        }

        let sourceRequest = NSFetchRequest<NSManagedObject>(entityName: "InvestmentSource")
        sourceRequest.resultType = .countResultType
        let sourceCount = (try? context.count(for: sourceRequest)) ?? 0
        return sourceCount > 0
    }

    private static func resolveStoreURL() -> URL {
        let defaultURL = NSPersistentContainer.defaultDirectoryURL()
            .appendingPathComponent("PortfolioJournal.sqlite")

        guard appGroupEnabled,
              let appGroupURL = FileManager.default
                .containerURL(forSecurityApplicationGroupIdentifier: appGroupIdentifier)?
                .appendingPathComponent("PortfolioJournal.sqlite") else {
            if appGroupEnabled {
                print("App Group unavailable; using default store URL: \(defaultURL)")
            }
            return defaultURL
        }

        migrateStoreIfNeeded(from: defaultURL, to: appGroupURL)
        return appGroupURL
    }

    @Published private(set) var isLoaded = false

    lazy var persistentContainer: NSPersistentContainer = {
        let container: NSPersistentContainer
        if Self.cloudKitEnabled {
            container = NSPersistentCloudKitContainer(name: "PortfolioJournal")
        } else {
            container = NSPersistentContainer(name: "PortfolioJournal")
        }

        // App Group store URL for sharing with widgets. Fall back if not entitled.
        let storeURL = Self.resolveStoreURL()

        let description = NSPersistentStoreDescription(url: storeURL)
        description.shouldMigrateStoreAutomatically = true
        description.shouldInferMappingModelAutomatically = true

        if Self.cloudKitEnabled {
            description.cloudKitContainerOptions = NSPersistentCloudKitContainerOptions(
                containerIdentifier: Self.cloudKitContainerIdentifier
            )
            description.setOption(true as NSNumber, forKey: NSPersistentHistoryTrackingKey)
            description.setOption(true as NSNumber, forKey: NSPersistentStoreRemoteChangeNotificationPostOptionKey)
        }

        container.persistentStoreDescriptions = [description]

        container.loadPersistentStores { [weak self] description, error in
            if let error = error as NSError? {
                // In production, handle this error appropriately
                print("Core Data failed to load: \(error), \(error.userInfo)")
            } else {
                print("Core Data loaded successfully at: \(description.url?.path ?? "unknown")")
            }
            DispatchQueue.main.async {
                self?.isLoaded = true
            }
        }

        // Merge policy - remote changes win
        container.viewContext.automaticallyMergesChangesFromParent = true
        container.viewContext.mergePolicy = NSMergeByPropertyObjectTrumpMergePolicy

        // Performance optimizations
        container.viewContext.undoManager = nil
        container.viewContext.shouldDeleteInaccessibleFaults = true

        if Self.cloudKitEnabled {
            NotificationCenter.default.addObserver(
                self,
                selector: #selector(processRemoteChanges),
                name: .NSPersistentStoreRemoteChange,
                object: container.persistentStoreCoordinator
            )
        }

        return container
    }()

    var viewContext: NSManagedObjectContext {
        return persistentContainer.viewContext
    }

    private init() {}

    // MARK: - Save Context

    func save() {
        let context = viewContext
        guard context.hasChanges else { return }

        do {
            try context.save()
        } catch {
            let nsError = error as NSError
            print("Core Data save error: \(nsError), \(nsError.userInfo)")
        }
    }

    // MARK: - Background Context

    func newBackgroundContext() -> NSManagedObjectContext {
        let context = persistentContainer.newBackgroundContext()
        context.mergePolicy = NSMergeByPropertyObjectTrumpMergePolicy
        context.undoManager = nil
        return context
    }

    func performBackgroundTask(_ block: @escaping (NSManagedObjectContext) -> Void) {
        persistentContainer.performBackgroundTask(block)
    }

    // MARK: - Remote Change Handling

    @objc private func processRemoteChanges(_ notification: Notification) {
        // Process remote changes on main context
        DispatchQueue.main.async { [weak self] in
            self?.objectWillChange.send()
            // Ensure changes are persisted to disk before refreshing widget
            self?.save()
            self?.refreshWidgetData()
        }
    }

    // MARK: - Widget Data Refresh

    func refreshWidgetData() {
        if #available(iOS 14.0, *) {
            if Self.appGroupEnabled {
                checkpointWAL()
                WidgetCenter.shared.reloadAllTimelines()
            }
        }
    }

    /// Forces SQLite to checkpoint the WAL file, ensuring all changes are written to the main database file.
    /// This is necessary for the widget to see changes, as it opens the database read-only.
    private func checkpointWAL() {
        if viewContext.hasChanges {
            try? viewContext.save()
        }

        let coordinator = persistentContainer.persistentStoreCoordinator
        coordinator.performAndWait {
            viewContext.refreshAllObjects()
        }
    }
}

// MARK: - Shared Container for Widgets

extension CoreDataStack {
    static var sharedStoreURL: URL? {
        guard appGroupEnabled else {
            let fallbackURL = NSPersistentContainer.defaultDirectoryURL()
                .appendingPathComponent("PortfolioJournal.sqlite")
            return fallbackURL
        }

        if let appGroupURL = FileManager.default
            .containerURL(forSecurityApplicationGroupIdentifier: appGroupIdentifier)?
            .appendingPathComponent("PortfolioJournal.sqlite") {
            return appGroupURL
        }

        let fallbackURL = NSPersistentContainer.defaultDirectoryURL()
            .appendingPathComponent("PortfolioJournal.sqlite")
        print("App Group unavailable for widgets; using default store URL: \(fallbackURL)")
        return fallbackURL
    }

    /// Creates a lightweight Core Data stack for widgets (read-only)
    static func createWidgetContainer() -> NSPersistentContainer {
        let container = NSPersistentContainer(name: "PortfolioJournal")

        guard let storeURL = sharedStoreURL else {
            fatalError("Unable to get shared store URL")
        }

        let description = NSPersistentStoreDescription(url: storeURL)
        description.isReadOnly = true
        description.shouldMigrateStoreAutomatically = true
        description.shouldInferMappingModelAutomatically = true

        container.persistentStoreDescriptions = [description]

        container.loadPersistentStores { _, error in
            if let error = error {
                print("Widget Core Data failed: \(error)")
            }
        }

        return container
    }
}