InvestmentTrackerApp/InvestmentTracker/Repositories/SnapshotRepository.swift

import Foundation
import CoreData
import Combine

class SnapshotRepository: ObservableObject {
    private let context: NSManagedObjectContext

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

    init(context: NSManagedObjectContext = CoreDataStack.shared.viewContext) {
        self.context = context
    }

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

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

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

        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
    ) {
        if let date = date {
            snapshot.date = date
        }
        if let value = value {
            snapshot.value = NSDecimalNumber(decimal: value)
        }
        if let contribution = contribution {
            snapshot.contribution = NSDecimalNumber(decimal: contribution)
        }
        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 {
            let cutoffDate = Calendar.current.date(
                byAdding: .month,
                value: -months,
                to: Date()
            ) ?? Date()

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

        return snapshots
    }

    // MARK: - Historical Data

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

        var monthlyData: [(date: Date, value: Decimal)] = []
        var processedMonths: Set<String> = []

        let formatter = DateFormatter()
        formatter.dateFormat = "yyyy-MM"

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

        // Group by date
        var dateValues: [Date: Decimal] = [:]

        for snapshot in allSnapshots {
            let startOfDay = Calendar.current.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()
        } catch {
            print("Failed to save context: \(error)")
        }
    }
}