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: - 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.
    func cleanupDuplicateObjects() {
        let context = viewContext
        context.performAndWait {
            // Clean up duplicate Goals
            cleanupDuplicates(entityName: "Goal", idKey: "id", context: context)
            // Clean up duplicate Accounts (already handled in AccountRepository but added here for safety)
            cleanupDuplicates(entityName: "Account", idKey: "id", context: context)
            // Clean up duplicate InvestmentSources
            cleanupDuplicates(entityName: "InvestmentSource", idKey: "id", context: context)
            // Clean up duplicate Categories
            cleanupDuplicates(entityName: "Category", idKey: "id", context: context)

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

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

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

        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("Cleaning up \(objectsToDelete.count) duplicate \(entityName) objects")
            for object in objectsToDelete {
                context.delete(object)
            }
        }
    }

    // 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,
            // which now contains the changes CloudKit just synced from other devices.
            // This triggers NSManagedObjectContextObjectsDidChange so all repositories
            // automatically re-fetch and push fresh data to ViewModels.
            self.viewContext.refreshAllObjects()
            self.objectWillChange.send()
            self.save()
            self.refreshWidgetData()
        }
    }

    /// 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
    }
}