InvestmentTrackerApp/PortfolioJournal/Models/CoreDataStack.swift

import CoreData
import CloudKit
import Combine
import WidgetKit

extension Notification.Name {
    static let cloudKitForceReload = Notification.Name("cloudKitForceReload")
}

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
    @Published private(set) var lastImportDate: Date?
    @Published private(set) var lastExportDate: Date?
    @Published private(set) var isSyncing = false
    @Published private(set) var lastSyncError: String?

    var lastSyncDate: Date? { lastImportDate }

    var localSourceCount: Int {
        let request = NSFetchRequest<NSManagedObject>(entityName: "InvestmentSource")
        return (try? viewContext.count(for: request)) ?? 0
    }

    var localSnapshotCount: Int {
        let request = NSFetchRequest<NSManagedObject>(entityName: "Snapshot")
        return (try? viewContext.count(for: request)) ?? 0
    }

    private init() {
        // Register CloudKit event observer BEFORE the container is created so we
        // never miss an import/export event that fires during store loading.
        if Self.cloudKitEnabled {
            NotificationCenter.default.addObserver(
                self,
                selector: #selector(cloudKitEventChanged(_:)),
                name: NSPersistentCloudKitContainer.eventChangedNotification,
                object: nil
            )
        }
    }

    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

        // Always enable history tracking so data created before CloudKit was enabled
        // is visible to NSPersistentCloudKitContainer when sync is later turned on.
        description.setOption(true as NSNumber, forKey: NSPersistentHistoryTrackingKey)
        description.setOption(true as NSNumber, forKey: NSPersistentStoreRemoteChangeNotificationPostOptionKey)

        if Self.cloudKitEnabled {
            description.cloudKitContainerOptions = NSPersistentCloudKitContainerOptions(
                containerIdentifier: Self.cloudKitContainerIdentifier
            )
        }

        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
    }

    // MARK: - Cleanup Duplicates

    /// Removes duplicate objects that have the same UUID, keeping only the oldest one.
    /// This fixes data corruption from race conditions during object creation, and also
    /// handles the case where a CloudKit first-time sync imports records that already exist
    /// locally (because they were created before CloudKit was enabled).
    @discardableResult
    func cleanupDuplicateObjects() -> Int {
        var totalRemoved = 0
        let context = viewContext
        context.performAndWait {
            // Clean up Snapshots first — before cascade rules from InvestmentSource fire,
            // so we deduplicate by UUID and not rely solely on cascade.
            totalRemoved += cleanupDuplicates(entityName: "Snapshot", idKey: "id", context: context)
            totalRemoved += cleanupDuplicates(entityName: "Goal", idKey: "id", context: context)
            totalRemoved += cleanupDuplicates(entityName: "Account", idKey: "id", context: context)
            totalRemoved += cleanupDuplicates(entityName: "InvestmentSource", idKey: "id", context: context)
            totalRemoved += cleanupDuplicates(entityName: "Category", idKey: "id", context: context)

            if context.hasChanges {
                try? context.save()
            }
        }
        return totalRemoved
    }

    @discardableResult
    private func cleanupDuplicates(entityName: String, idKey: String, context: NSManagedObjectContext) -> Int {
        let request = NSFetchRequest<NSManagedObject>(entityName: entityName)
        request.sortDescriptors = [NSSortDescriptor(key: "createdAt", ascending: true)]

        guard let objects = try? context.fetch(request) else { return 0 }

        var seenIds = Set<UUID>()
        var objectsToDelete: [NSManagedObject] = []

        for object in objects {
            guard let objectId = object.value(forKey: idKey) as? UUID else { continue }
            if seenIds.contains(objectId) {
                objectsToDelete.append(object)
            } else {
                seenIds.insert(objectId)
            }
        }

        if !objectsToDelete.isEmpty {
            print("[Dedup] Removing \(objectsToDelete.count) duplicate \(entityName) objects")
            for object in objectsToDelete { context.delete(object) }
        }
        return objectsToDelete.count
    }

    // 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) {
        DispatchQueue.main.async { [weak self] in
            guard let self else { return }
            // Force viewContext to re-read all objects from the persistent store.
            self.viewContext.refreshAllObjects()
            // Remove any duplicates that CloudKit import may have introduced.
            // This handles the first-time sync case where records existed locally before
            // CloudKit was enabled: the initial export+import creates duplicate objects.
            let removed = self.cleanupDuplicateObjects()
            if removed > 0 {
                print("[RemoteChanges] Removed \(removed) duplicate objects after CloudKit import")
            }
            // Notify repositories to re-fetch unconditionally.
            NotificationCenter.default.post(name: .cloudKitForceReload, object: nil)
            self.objectWillChange.send()
            self.save()
            self.refreshWidgetData()
        }
    }

    @objc private func cloudKitEventChanged(_ notification: Notification) {
        guard let event = notification.userInfo?[NSPersistentCloudKitContainer.eventNotificationUserInfoKey]
                as? NSPersistentCloudKitContainer.Event else { return }
        DispatchQueue.main.async { [weak self] in
            guard let self else { return }
            let isActive = event.endDate == nil
            self.isSyncing = isActive
            if !isActive {
                if event.succeeded, let endDate = event.endDate {
                    self.lastSyncError = nil
                    switch event.type {
                    case .import: self.lastImportDate = endDate
                    case .export: self.lastExportDate = endDate
                    default: break
                    }
                } else if let error = event.error {
                    let typeLabel = event.type == .import ? "import" : event.type == .export ? "export" : "setup"
                    self.lastSyncError = "\(typeLabel): \(Self.describeError(error))"
                    print("CloudKit \(typeLabel) full error:\n\(error)\nuserInfo: \((error as NSError).userInfo)")
                }
            }
        }
    }

    private static func describeError(_ error: Error, depth: Int = 0) -> String {
        guard depth < 3 else { return "…" }
        let ns = error as NSError
        var parts: [String] = ["\(ns.domain)(\(ns.code))"]
        // At depth 0, always include the localised description so the user sees
        // a human-readable message even when userInfo has no other keys.
        if depth == 0, let msg = ns.userInfo[NSLocalizedDescriptionKey] as? String {
            parts.append(msg)
        }
        for (key, value) in ns.userInfo {
            let k = "\(key)"
            if k == NSLocalizedDescriptionKey || k == "NSLocalizedDescription" { continue }
            if let nestedError = value as? Error {
                parts.append("\(k):\(describeError(nestedError, depth: depth + 1))")
            } else if let dict = value as? [AnyHashable: Any], !dict.isEmpty {
                let pairs = dict.prefix(3).map { kk, vv -> String in
                    if let e = vv as? Error { return "\(kk)→\(describeError(e, depth: depth + 1))" }
                    return "\(kk)=\(vv)"
                }
                parts.append("\(k){\(pairs.joined(separator: ", "))}")
            } else if !(value is [AnyHashable: Any]) {
                parts.append("\(k)=\(value)")
            }
        }
        return parts.joined(separator: "\n")
    }

    func forceReload() {
        viewContext.perform { [weak self] in
            self?.viewContext.refreshAllObjects()
            DispatchQueue.main.async {
                NotificationCenter.default.post(name: .cloudKitForceReload, object: nil)
            }
        }
    }

    /// Forces NSPersistentCloudKitContainer to export all local data to iCloud.
    ///
    /// Strategy: advance `createdAt` by 1 ms for every record.  This guarantees
    /// a *real* value change that the SQLite persistent store will write as a
    /// persistent-history transaction — which is what NSPersistentCloudKitContainer
    /// needs to discover records and enqueue them for CloudKit export.
    ///
    /// Setting a property to the *same* value may be silently discarded by the
    /// SQLite layer (no SQL UPDATE issued → no history entry → nothing to export).
    func forceExportToiCloud(completion: @escaping (Int) -> Void) {
        guard Self.cloudKitEnabled else { completion(0); return }
        let entities = ["Account", "Category", "InvestmentSource", "Snapshot", "Goal"]
        let context = newBackgroundContext()
        context.perform { [weak self] in
            var totalTouched = 0
            for entityName in entities {
                let request = NSFetchRequest<NSManagedObject>(entityName: entityName)
                request.fetchBatchSize = 50
                guard let objects = try? context.fetch(request) else { continue }
                for object in objects {
                    // Advance createdAt by 1 ms → always a genuine value change.
                    let t = (object.value(forKey: "createdAt") as? Date) ?? Date()
                    object.setValue(t.addingTimeInterval(0.001), forKey: "createdAt")
                    totalTouched += 1
                }
            }
            if context.hasChanges {
                try? context.save()
            }
            DispatchQueue.main.async {
                completion(totalTouched)
            }
        }
    }

    /// Forces an immediate re-read of all Core Data objects from the persistent store.
    /// Call this when the app returns to the foreground so any iCloud changes made on
    /// other devices (while this device was inactive) are picked up right away.
    func refreshFromCloudKit() {
        guard Self.cloudKitEnabled else { return }
        viewContext.perform { [weak self] in
            self?.viewContext.refreshAllObjects()
        }
    }

    // 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() {
        viewContext.performAndWait {
            if viewContext.hasChanges {
                try? viewContext.save()
            }
            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
    }
}