InvestmentTrackerApp/PortfolioJournal/Repositories/SnapshotRepository.swift

import Foundation
import CoreData
import Combine

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] = []

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

        save()
        return snapshot
    }

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

    private func setupNotificationObserver() {
        NotificationCenter.default.addObserver(
            self,
            selector: #selector(contextDidChange),
            name: .NSManagedObjectContextObjectsDidChange,
            object: context
        )
    }

    @objc private func contextDidChange(_ notification: Notification) {
        invalidateCache()
    }
}