InvestmentTrackerApp/PortfolioJournal/Repositories/SnapshotRepository.swift

import Foundation
import CoreData
import Combine

@MainActor
class SnapshotRepository: ObservableObject {
    private let context: NSManagedObjectContext
    private let cache = NSCache<NSString, NSArray>()
    @Published private(set) var cacheVersion: Int = 0

    @Published private(set) var snapshots: [Snapshot] = []

    private var cancellables = Set<AnyCancellable>()

    // MARK: - Performance: Shared DateFormatter
    private static let monthYearFormatter: DateFormatter = {
        let formatter = DateFormatter()
        formatter.dateFormat = "yyyy-MM"
        return formatter
    }()

    // MARK: - Performance: Cached Calendar
    private static let calendar = Calendar.current

    init(context: NSManagedObjectContext = CoreDataStack.shared.viewContext) {
        self.context = context
        // Performance: Increase cache limits for better hit rate
        cache.countLimit = 12
        cache.totalCostLimit = 4_000_000 // 4 MB
        setupNotificationObserver()
    }

    // MARK: - Fetch

    func fetchSnapshots(for source: InvestmentSource) -> [Snapshot] {
        let request: NSFetchRequest<Snapshot> = Snapshot.fetchRequest()
        request.predicate = NSPredicate(format: "source == %@", source)
        request.sortDescriptors = [
            NSSortDescriptor(keyPath: \Snapshot.date, ascending: false)
        ]
        request.fetchBatchSize = 200

        return (try? context.fetch(request)) ?? []
    }

    func fetchSnapshot(by id: UUID) -> Snapshot? {
        let request: NSFetchRequest<Snapshot> = Snapshot.fetchRequest()
        request.predicate = NSPredicate(format: "id == %@", id as CVarArg)
        request.fetchLimit = 1
        return try? context.fetch(request).first
    }

    func fetchAllSnapshots() -> [Snapshot] {
        let request: NSFetchRequest<Snapshot> = Snapshot.fetchRequest()
        request.sortDescriptors = [
            NSSortDescriptor(keyPath: \Snapshot.date, ascending: false)
        ]
        request.fetchBatchSize = 200

        return (try? context.fetch(request)) ?? []
    }

    func fetchSnapshots(for account: Account) -> [Snapshot] {
        let request: NSFetchRequest<Snapshot> = Snapshot.fetchRequest()
        request.predicate = NSPredicate(format: "source.account == %@", account)
        request.sortDescriptors = [
            NSSortDescriptor(keyPath: \Snapshot.date, ascending: false)
        ]
        request.fetchBatchSize = 200

        return (try? context.fetch(request)) ?? []
    }

    func fetchSnapshots(from startDate: Date, to endDate: Date) -> [Snapshot] {
        let request: NSFetchRequest<Snapshot> = Snapshot.fetchRequest()
        request.predicate = NSPredicate(
            format: "date >= %@ AND date <= %@",
            startDate as NSDate,
            endDate as NSDate
        )
        request.sortDescriptors = [
            NSSortDescriptor(keyPath: \Snapshot.date, ascending: true)
        ]
        request.fetchBatchSize = 200

        return (try? context.fetch(request)) ?? []
    }

    func fetchLatestSnapshots() -> [Snapshot] {
        // Get the latest snapshot for each source
        let request: NSFetchRequest<InvestmentSource> = InvestmentSource.fetchRequest()
        let sources = (try? context.fetch(request)) ?? []

        return sources.compactMap { $0.latestSnapshot }
    }

    // MARK: - Create

    @discardableResult
    func createSnapshot(
        for source: InvestmentSource,
        date: Date = Date(),
        value: Decimal,
        contribution: Decimal? = nil,
        notes: String? = nil,
        batch: Bool = false
    ) -> Snapshot {
        let snapshot = Snapshot(context: context)
        snapshot.source = source
        snapshot.date = date
        snapshot.value = NSDecimalNumber(decimal: value)
        if let contribution = contribution {
            snapshot.contribution = NSDecimalNumber(decimal: contribution)
        }
        snapshot.notes = notes

        // In batch mode (CSV import), skip the per-row save + widget refresh — the
        // caller saves the whole context once at the end. Saving per row triggered a
        // main-thread WAL checkpoint each time → ANR on large imports.
        if !batch {
            save()
        }
        return snapshot
    }

    // MARK: - Gap Auto-Fill (estimated snapshots)

    /// Creates ESTIMATED snapshots for the missing months of an internal gap,
    /// linearly interpolated between the source's value at `from` and at `to`.
    /// They are marked `isEstimated` so they read as clearly distinct from real
    /// data and can be removed in bulk. Idempotent (skips months that already have
    /// a snapshot). Returns how many were created.
    @discardableResult
    func autoFillGap(source: InvestmentSource, from: Date, to: Date, missingMonthDates: [Date]) -> Int {
        let snaps = source.snapshotsArray
        func value(atMonth month: Date) -> Decimal? {
            snaps.filter { $0.date.startOfMonth == month.startOfMonth }
                 .max(by: { $0.date < $1.date })?.decimalValue
        }
        guard let v0 = value(atMonth: from),
              let v1 = value(atMonth: to),
              !missingMonthDates.isEmpty else { return 0 }

        let steps = missingMonthDates.count + 1   // segments between v0 and v1
        var created = 0
        for (i, month) in missingMonthDates.enumerated() {
            let monthStart = month.startOfMonth
            if snaps.contains(where: { $0.date.startOfMonth == monthStart }) { continue }
            let fraction = Decimal(i + 1) / Decimal(steps)
            let interpolated = v0 + (v1 - v0) * fraction
            let snapshot = Snapshot(context: context)
            snapshot.source = source
            snapshot.date = monthStart
            snapshot.value = NSDecimalNumber(decimal: interpolated)
            snapshot.isEstimated = true
            snapshot.notes = String(localized: "gaps_estimated_note")
            created += 1
        }
        if created > 0 { save() }
        return created
    }

    /// Deletes all estimated (auto-filled) snapshots, optionally scoped to one
    /// source. Returns how many were removed.
    @discardableResult
    func deleteEstimatedSnapshots(for source: InvestmentSource? = nil) -> Int {
        let request = NSFetchRequest<Snapshot>(entityName: "Snapshot")
        request.predicate = NSPredicate(format: "isEstimated == YES")
        guard let estimated = try? context.fetch(request) else { return 0 }
        let targets = source == nil ? estimated : estimated.filter { $0.source?.id == source?.id }
        for snapshot in targets { context.delete(snapshot) }
        if !targets.isEmpty { save() }
        return targets.count
    }

    // MARK: - Update

    func updateSnapshot(
        _ snapshot: Snapshot,
        date: Date? = nil,
        value: Decimal? = nil,
        contribution: Decimal? = nil,
        notes: String? = nil,
        clearContribution: Bool = false,
        clearNotes: Bool = false
    ) {
        if let date = date {
            snapshot.date = date
        }
        if let value = value {
            snapshot.value = NSDecimalNumber(decimal: value)
        }
        if clearContribution {
            snapshot.contribution = nil
        } else if let contribution = contribution {
            snapshot.contribution = NSDecimalNumber(decimal: contribution)
        }
        if clearNotes {
            snapshot.notes = nil
        } else if let notes = notes {
            snapshot.notes = notes
        }

        save()
    }

    // MARK: - Contribution Propagation

    enum ContributionPropagation {
        case forward   // snapshots after the pivot
        case backward  // snapshots before the pivot
        case all       // every other snapshot of the source
    }

    /// Applies `amount` as the contribution of other snapshots of the same source,
    /// relative to `snapshot`'s date, according to `direction`. The pivot snapshot
    /// itself is left untouched (it was already saved with its own value).
    func propagateContribution(
        _ amount: Decimal,
        from snapshot: Snapshot,
        direction: ContributionPropagation
    ) {
        guard let source = snapshot.source,
              let set = source.snapshots as? Set<Snapshot> else { return }
        let pivot = snapshot.date
        let targets = set.filter { other in
            guard other.objectID != snapshot.objectID else { return false }
            switch direction {
            case .forward: return other.date > pivot
            case .backward: return other.date < pivot
            case .all: return true
            }
        }
        let value = NSDecimalNumber(decimal: amount)
        for target in targets {
            target.contribution = value
        }
        save()
    }

    // MARK: - Delete

    func deleteSnapshot(_ snapshot: Snapshot) {
        context.delete(snapshot)
        save()
    }

    func deleteSnapshots(_ snapshots: [Snapshot]) {
        snapshots.forEach { context.delete($0) }
        save()
    }

    func deleteSnapshot(at offsets: IndexSet, from snapshots: [Snapshot]) {
        for index in offsets {
            guard index < snapshots.count else { continue }
            context.delete(snapshots[index])
        }
        save()
    }

    // MARK: - Filtered Fetch (Freemium)

    func fetchSnapshots(
        for source: InvestmentSource,
        limitedToMonths months: Int?
    ) -> [Snapshot] {
        var snapshots = fetchSnapshots(for: source)

        if let months = months {
            // Performance: Use cached calendar
            let cutoffDate = Self.calendar.date(
                byAdding: .month,
                value: -months,
                to: Date()
            ) ?? Date()

            snapshots = snapshots.filter { $0.date >= cutoffDate }
        }

        return snapshots
    }

    func fetchSnapshots(
        for sourceIds: [UUID],
        months: Int? = nil
    ) -> [Snapshot] {
        guard !sourceIds.isEmpty else { return [] }

        // Performance: Use cached calendar
        let cutoffDate: Date? = {
            guard let months = months else { return nil }
            return Self.calendar.date(byAdding: .month, value: -months, to: Date())
        }()

        let key = cacheKey(
            sourceIds: sourceIds,
            months: months
        )

        if let cached = cache.object(forKey: key) as? [Snapshot] {
            return cached
        }

        let request: NSFetchRequest<Snapshot> = Snapshot.fetchRequest()
        var predicates: [NSPredicate] = [
            NSPredicate(format: "source.id IN %@", sourceIds)
        ]

        if let cutoffDate {
            predicates.append(NSPredicate(format: "date >= %@", cutoffDate as NSDate))
        }

        request.predicate = NSCompoundPredicate(andPredicateWithSubpredicates: predicates)
        request.sortDescriptors = [
            NSSortDescriptor(keyPath: \Snapshot.date, ascending: false)
        ]
        request.fetchBatchSize = 300

        let fetched = (try? context.fetch(request)) ?? []
        cache.setObject(fetched as NSArray, forKey: key, cost: fetched.count)
        return fetched
    }

    // MARK: - Historical Data

    func getMonthlyValues(for source: InvestmentSource) -> [(date: Date, value: Decimal)] {
        let snapshots = fetchSnapshots(for: source).reversed()

        var monthlyData: [(date: Date, value: Decimal)] = []
        monthlyData.reserveCapacity(min(snapshots.count, 60))
        var processedMonths: Set<String> = []
        processedMonths.reserveCapacity(60)

        // Performance: Use shared formatter
        let formatter = Self.monthYearFormatter

        for snapshot in snapshots {
            let monthKey = formatter.string(from: snapshot.date)
            if !processedMonths.contains(monthKey) {
                processedMonths.insert(monthKey)
                monthlyData.append((date: snapshot.date, value: snapshot.decimalValue))
            }
        }

        return monthlyData
    }

    func getPortfolioHistory() -> [(date: Date, totalValue: Decimal)] {
        let allSnapshots = fetchAllSnapshots()

        // Performance: Pre-allocate dictionary
        var dateValues: [Date: Decimal] = [:]
        dateValues.reserveCapacity(min(allSnapshots.count, 365))

        // Performance: Use cached calendar
        let calendar = Self.calendar

        for snapshot in allSnapshots {
            let startOfDay = calendar.startOfDay(for: snapshot.date)
            dateValues[startOfDay, default: Decimal.zero] += snapshot.decimalValue
        }

        return dateValues
            .map { (date: $0.key, totalValue: $0.value) }
            .sorted { $0.date < $1.date }
    }

    // MARK: - Save

    private func save() {
        guard context.hasChanges else { return }
        do {
            try context.save()
            invalidateCache()
            CoreDataStack.shared.refreshWidgetData()
        } catch {
            print("Failed to save context: \(error)")
        }
    }

    // MARK: - Cache

    private func cacheKey(
        sourceIds: [UUID],
        months: Int?
    ) -> NSString {
        let sortedIds = sourceIds.sorted().map { $0.uuidString }.joined(separator: ",")
        let monthsPart = months.map(String.init) ?? "all"
        return NSString(string: "\(cacheVersion)|\(monthsPart)|\(sortedIds)")
    }

    private func invalidateCache() {
        cache.removeAllObjects()
        cacheVersion &+= 1
    }

    /// Clear cache on memory pressure - call from AppDelegate's didReceiveMemoryWarning
    func clearCacheOnMemoryPressure() {
        cache.removeAllObjects()
    }

    private func setupNotificationObserver() {
        NotificationCenter.default.publisher(for: .NSManagedObjectContextObjectsDidChange, object: context)
            .receive(on: DispatchQueue.main)
            .sink { [weak self] _ in
                self?.invalidateCache()
            }
            .store(in: &cancellables)
    }
}