InvestmentTrackerApp/PortfolioJournal/ViewModels/ChartsViewModel.swift

import Foundation
import Combine
import CoreData

@MainActor
class ChartsViewModel: ObservableObject {
    // MARK: - Chart Types

    enum ChartType: String, CaseIterable, Identifiable {
        case evolution = "Evolution"
        case allocation = "Allocation"
        case performance = "Performance"
        case contributions = "Contributions"
        case rollingReturn = "Rolling 12M"
        case riskReturn = "Risk vs Return"
        case cashflow = "Net vs Contributions"
        case drawdown = "Drawdown"
        case volatility = "Volatility"
        case prediction = "Prediction"

        var id: String { rawValue }

        var icon: String {
            switch self {
            case .evolution: return "chart.line.uptrend.xyaxis"
            case .allocation: return "chart.pie.fill"
            case .performance: return "chart.bar.fill"
            case .contributions: return "tray.and.arrow.down.fill"
            case .rollingReturn: return "arrow.triangle.2.circlepath"
            case .riskReturn: return "dot.square"
            case .cashflow: return "chart.bar.xaxis"
            case .drawdown: return "arrow.down.right.circle"
            case .volatility: return "waveform.path.ecg"
            case .prediction: return "wand.and.stars"
            }
        }

        var isPremium: Bool {
            switch self {
            case .evolution:
                return false
            case .allocation, .performance, .contributions, .rollingReturn, .riskReturn, .cashflow,
                 .drawdown, .volatility, .prediction:
                return true
            }
        }

        var description: String {
            switch self {
            case .evolution:
                return "Track your portfolio value over time"
            case .allocation:
                return "See how your investments are distributed"
            case .performance:
                return "Compare returns across categories"
            case .contributions:
                return "Review monthly inflows over time"
            case .rollingReturn:
                return "See rolling 12-month performance"
            case .riskReturn:
                return "Compare volatility vs return"
            case .cashflow:
                return "Compare growth vs contributions"
            case .drawdown:
                return "Analyze declines from peak values"
            case .volatility:
                return "Understand investment risk levels"
            case .prediction:
                return "View 12-month forecasts"
            }
        }
    }

    func availableChartTypes(calmModeEnabled: Bool) -> [ChartType] {
        let types: [ChartType] = calmModeEnabled
            ? [.evolution, .allocation, .performance, .contributions]
            : ChartType.allCases
        if !types.contains(selectedChartType) {
            selectedChartType = .evolution
        }
        return types
    }

    // MARK: - Published Properties

    @Published var selectedChartType: ChartType = .evolution
    @Published var selectedCategory: Category?
    @Published var selectedSource: InvestmentSource?
    @Published var selectedSourceIds: Set<UUID> = []
    @Published var selectedTimeRange: TimeRange = .year
    @Published var selectedAccount: Account?
    @Published var showAllAccounts = true
    @Published var selectedBreakdown: BreakdownMode = .category

    @Published var evolutionData: [(date: Date, value: Decimal)] = []
    @Published var categoryEvolutionData: [CategoryEvolutionPoint] = []
    @Published var allocationData: [(category: String, value: Decimal, color: String)] = []
    @Published var performanceData: [(category: String, cagr: Double, color: String)] = []
    @Published var contributionsData: [(date: Date, amount: Decimal)] = []
    @Published var rollingReturnData: [(date: Date, value: Double)] = []
    @Published var riskReturnData: [(category: String, cagr: Double, volatility: Double, color: String)] = []
    @Published var cashflowData: [(date: Date, contributions: Decimal, netPerformance: Decimal)] = []
    @Published var drawdownData: [(date: Date, drawdown: Double)] = []
    @Published var volatilityData: [(date: Date, volatility: Double)] = []
    @Published var predictionData: [Prediction] = []
    @Published var allocationEvolutionData: [(date: Date, category: String, percentage: Double, color: String)] = []

    @Published var isLoading = false
    @Published var showingPaywall = false
    @Published private var predictionMonthsAhead = 12

    // MARK: - Time Range

    enum TimeRange: String, CaseIterable, Identifiable {
        case month = "1M"
        case quarter = "3M"
        case halfYear = "6M"
        case year = "12M"
        case yearToDate = "YTD"
        case all = "All"

        var id: String { rawValue }

        var months: Int? {
            switch self {
            case .month: return 1
            case .quarter: return 3
            case .halfYear: return 6
            case .year: return 12
            case .yearToDate: return nil
            case .all: return nil
            }
        }

        func startDate(referenceDate: Date = Date()) -> Date? {
            switch self {
            case .month, .quarter, .halfYear, .year:
                guard let months else { return nil }
                return referenceDate.adding(months: -months).startOfDay
            case .yearToDate:
                return referenceDate.startOfYear
            case .all:
                return nil
            }
        }
    }

    enum BreakdownMode: String, CaseIterable, Identifiable {
        case category = "By Category"
        case source = "By Source"

        var id: String { rawValue }
    }

    static func supportsAllocationTargets(for breakdown: BreakdownMode) -> Bool {
        breakdown == .category
    }

    // MARK: - Dependencies

    private let sourceRepository: InvestmentSourceRepository
    private let categoryRepository: CategoryRepository
    private let snapshotRepository: SnapshotRepository
    private let calculationService: CalculationService
    private let predictionEngine: PredictionEngine
    private let freemiumValidator: FreemiumValidator
    private let maxHistoryMonths = 60
    private let maxStackedCategories = 6
    private let maxChartPoints = 500
    private var cancellables = Set<AnyCancellable>()
    private var allCategories: [Category] {
        categoryRepository.categories
    }

    // MARK: - Performance: Caching and State
    private var lastChartType: ChartType?
    private var lastTimeRange: TimeRange?
    private var lastCategoryId: UUID?
    private var lastSourceId: UUID?
    private var lastSourceIds: Set<UUID> = []
    private var lastAccountId: UUID?
    private var lastShowAllAccounts: Bool = true
    private var lastBreakdown: BreakdownMode = .category
    private var cachedSnapshots: [Snapshot]?
    private var isUpdateInProgress = false

    // MARK: - Initialization

    init(
        sourceRepository: InvestmentSourceRepository? = nil,
        categoryRepository: CategoryRepository? = nil,
        snapshotRepository: SnapshotRepository? = nil,
        calculationService: CalculationService? = nil,
        predictionEngine: PredictionEngine? = nil,
        iapService: IAPService
    ) {
        self.sourceRepository = sourceRepository ?? InvestmentSourceRepository()
        self.categoryRepository = categoryRepository ?? CategoryRepository()
        self.snapshotRepository = snapshotRepository ?? SnapshotRepository()
        self.calculationService = calculationService ?? .shared
        self.predictionEngine = predictionEngine ?? .shared
        self.freemiumValidator = FreemiumValidator(iapService: iapService)

        setupObservers()
        loadData()
    }

    // MARK: - Setup

    private func setupObservers() {
        // Performance: Combine all selection changes into a single debounced stream
        // This prevents multiple rapid updates when switching between views
        Publishers.CombineLatest4($selectedChartType, $selectedCategory, $selectedTimeRange, $selectedAccount)
            .combineLatest($selectedSource, $selectedBreakdown, $showAllAccounts)
            .combineLatest($selectedSourceIds)
            .debounce(for: .milliseconds(150), scheduler: DispatchQueue.main)
            .sink { [weak self] outer, sourceIds in
                guard let self else { return }
                let (combined, selectedSource, selectedBreakdown, showAll) = outer
                let (chartType, category, timeRange, _) = combined

                // Performance: Skip update if nothing meaningful changed
                let safeSelectedAccountId = self.safeSelectedAccountId
                let hasChanges = self.lastChartType != chartType ||
                    self.lastTimeRange != timeRange ||
                    self.lastCategoryId != category?.id ||
                    self.lastSourceId != selectedSource?.id ||
                    self.lastSourceIds != sourceIds ||
                    self.lastAccountId != safeSelectedAccountId ||
                    self.lastShowAllAccounts != showAll ||
                    self.lastBreakdown != selectedBreakdown

                if hasChanges {
                    self.lastChartType = chartType
                    self.lastTimeRange = timeRange
                    self.lastCategoryId = category?.id
                    self.lastSourceId = selectedSource?.id
                    self.lastSourceIds = sourceIds
                    self.lastAccountId = safeSelectedAccountId
                    self.lastShowAllAccounts = showAll
                    self.lastBreakdown = selectedBreakdown
                    self.cachedSnapshots = nil // Invalidate cache on meaningful changes
                    self.updateChartData(chartType: chartType, category: category, timeRange: timeRange)
                }
            }
            .store(in: &cancellables)
    }

    // MARK: - Data Loading

    func loadData() {
        updateChartData(
            chartType: selectedChartType,
            category: selectedCategory,
            timeRange: selectedTimeRange
        )

        FirebaseService.shared.logScreenView(screenName: "Charts")
    }

    func selectChart(_ chartType: ChartType) {
        if chartType.isPremium && !freemiumValidator.isPremium {
            showingPaywall = true
            FirebaseService.shared.logPaywallShown(trigger: "advanced_charts")
            return
        }

        selectedChartType = chartType
        let allowedRanges = availableTimeRanges(for: chartType)
        if !allowedRanges.contains(selectedTimeRange) {
            selectedTimeRange = allowedRanges.first ?? .year
        }
        FirebaseService.shared.logChartViewed(
            chartType: chartType.rawValue,
            isPremium: chartType.isPremium
        )
    }

    func availableTimeRanges(for chartType: ChartType) -> [TimeRange] {
        switch chartType {
        case .evolution, .performance:
            return [.all, .yearToDate, .year, .quarter]
        default:
            return [.month, .quarter, .halfYear, .year, .all]
        }
    }

    private func updateChartData(chartType: ChartType, category: Category?, timeRange: TimeRange) {
        // Performance: Prevent re-entrancy
        guard !isUpdateInProgress else { return }
        isUpdateInProgress = true
        isLoading = true

        defer {
            isLoading = false
            isUpdateInProgress = false
        }

        let sources: [InvestmentSource]
        if !selectedSourceIds.isEmpty {
            sources = sourceRepository.sources.filter { source in
                selectedSourceIds.contains(source.id) && shouldIncludeSource(source)
            }
        } else if let selectedSource {
            sources = sourceRepository.sources.filter { $0.id == selectedSource.id && shouldIncludeSource($0) }
        } else if let category = category {
            sources = sourceRepository.fetchSources(for: category).filter { shouldIncludeSource($0) }
        } else {
            sources = sourceRepository.sources.filter { shouldIncludeSource($0) }
        }

        let monthsLimit = timeRange.months ?? maxHistoryMonths
        let sourceIds = sources.compactMap { $0.id }

        // Performance: Reuse cached snapshots when possible
        var snapshots: [Snapshot]
        if let cached = cachedSnapshots {
            snapshots = cached
        } else {
            snapshots = snapshotRepository.fetchSnapshots(
                for: sourceIds,
                months: monthsLimit
            )
            snapshots = freemiumValidator.filterSnapshots(snapshots)
            cachedSnapshots = snapshots
        }

        if let cutoffDate = timeRange.startDate() {
            snapshots = snapshots.filter { $0.date >= cutoffDate }
        }

        let completedSnapshots = filterSnapshotsForCharts(
            sources: sources,
            snapshots: snapshots
        )

        // Performance: Only calculate data for the selected chart type
        switch chartType {
        case .evolution:
            calculateEvolutionData(from: completedSnapshots)
            let categoriesForChart = categoriesForStackedChart(
                sources: sources,
                selectedCategory: selectedCategory
            )
            calculateCategoryEvolutionData(from: completedSnapshots, categories: categoriesForChart)
        case .allocation:
            calculateAllocationData(for: sources, breakdown: selectedBreakdown)
            calculateAllocationEvolutionData(from: completedSnapshots)
        case .performance:
            let completedSnapshotsBySource = groupSnapshotsBySource(completedSnapshots)
            calculatePerformanceData(
                for: sources,
                snapshotsBySource: completedSnapshotsBySource,
                breakdown: selectedBreakdown
            )
        case .contributions:
            calculateContributionsData(from: completedSnapshots)
        case .rollingReturn:
            calculateRollingReturnData(from: completedSnapshots)
        case .riskReturn:
            let completedSnapshotsBySource = groupSnapshotsBySource(completedSnapshots)
            calculateRiskReturnData(for: sources, snapshotsBySource: completedSnapshotsBySource)
        case .cashflow:
            calculateCashflowData(from: completedSnapshots)
        case .drawdown:
            calculateDrawdownData(from: completedSnapshots)
        case .volatility:
            calculateVolatilityData(from: completedSnapshots)
        case .prediction:
            calculatePredictionData(from: completedSnapshots)
        }

        if let selected = selectedCategory,
           !availableCategories(for: chartType, sources: sources).contains(where: { $0.id == selected.id }) {
            selectedCategory = nil
        }
        if let selected = selectedSource,
           !availableSources(for: chartType, sources: sources).contains(where: { $0.id == selected.id }) {
            selectedSource = nil
        }
    }

    func availableCategories(
        for chartType: ChartType,
        sources: [InvestmentSource]? = nil
    ) -> [Category] {
        let relevantSources = sources ?? sourceRepository.sources.filter { shouldIncludeSource($0) }
        let categoriesWithData = Set(relevantSources.compactMap { $0.category?.id })
        let filtered = allCategories.filter { categoriesWithData.contains($0.id) }

        switch chartType {
        case .evolution, .prediction, .allocation, .performance:
            return filtered
        default:
            return []
        }
    }

    func availableSources(
        for chartType: ChartType,
        sources: [InvestmentSource]? = nil
    ) -> [InvestmentSource] {
        let relevantSources = sources ?? sourceRepository.sources.filter { shouldIncludeSource($0) }
        switch chartType {
        case .evolution, .allocation, .performance:
            return relevantSources.sorted { $0.name.localizedCaseInsensitiveCompare($1.name) == .orderedAscending }
        default:
            return []
        }
    }

    private func shouldIncludeSource(_ source: InvestmentSource) -> Bool {
        if showAllAccounts || selectedAccount == nil {
            return true
        }
        guard let selectedId = safeSelectedAccountId else { return true }
        return source.account?.id == selectedId
    }

    private var safeSelectedAccountId: UUID? {
        guard !showAllAccounts,
              let selected = selectedAccount,
              !selected.isDeleted else {
            return nil
        }
        return selected.id
    }

    /// Source-specific color palette (distinct from category colors, same order as Color.sourceColors)
    private static let sourceColorHexes: [String] = [
        "#6366F1", "#F97316", "#06B6D4", "#EF4444", "#84CC16", "#EC4899",
        "#14B8A6", "#F59E0B", "#8B5CF6", "#3B82F6", "#A855F7", "#10B981"
    ]

    /// Assigns a unique color to each source based on its sorted position among all sources.
    private func sourceColorMap(for sources: [InvestmentSource]) -> [UUID: String] {
        let sorted = sources.sorted { $0.name.localizedCaseInsensitiveCompare($1.name) == .orderedAscending }
        var map: [UUID: String] = [:]
        for (index, source) in sorted.enumerated() {
            map[source.id] = Self.sourceColorHexes[index % Self.sourceColorHexes.count]
        }
        return map
    }

    private func categoriesForStackedChart(
        sources: [InvestmentSource],
        selectedCategory: Category?
    ) -> [Category] {
        var totals: [UUID: Decimal] = [:]

        for source in sources {
            guard let categoryId = source.category?.id else { continue }
            totals[categoryId, default: 0] += source.latestValue
        }

        var topCategoryIds = Set(
            totals.sorted { $0.value > $1.value }
                .prefix(maxStackedCategories)
                .map { $0.key }
        )

        if let selectedCategory {
            topCategoryIds.insert(selectedCategory.id)
        }

        return categoryRepository.categories.filter { topCategoryIds.contains($0.id) }
    }

    private func downsampleSeries(
        _ data: [(date: Date, value: Decimal)],
        maxPoints: Int
    ) -> [(date: Date, value: Decimal)] {
        guard data.count > maxPoints, maxPoints > 0 else { return data }
        let bucketSize = max(1, Int(ceil(Double(data.count) / Double(maxPoints))))
        var sampled: [(date: Date, value: Decimal)] = []
        sampled.reserveCapacity(maxPoints)

        var index = 0
        while index < data.count {
            let end = min(index + bucketSize, data.count)
            let bucket = data[index..<end]
            if let last = bucket.last {
                sampled.append(last)
            }
            index += bucketSize
        }
        return sampled
    }

    private func downsampleDates(_ dates: [Date], maxPoints: Int) -> [Date] {
        guard dates.count > maxPoints, maxPoints > 0 else { return dates }
        let bucketSize = max(1, Int(ceil(Double(dates.count) / Double(maxPoints))))
        var sampled: [Date] = []
        sampled.reserveCapacity(maxPoints)

        var index = 0
        while index < dates.count {
            let end = min(index + bucketSize, dates.count)
            let bucket = dates[index..<end]
            if let last = bucket.last {
                sampled.append(last)
            }
            index += bucketSize
        }
        return sampled
    }

    // MARK: - Effective Month Mapping

    /// Maps a snapshot date to its effective check-in month.
    /// Snapshots on days 1–20 are attributed to the previous month's check-in.
    private func chartMonth(for snapshotDate: Date) -> DateComponents {
        let effective = MonthlyCheckInStore.effectiveMonth(for: snapshotDate, relativeTo: snapshotDate)
        return Calendar.current.dateComponents([.year, .month], from: effective)
    }

    private func chartMonthStart(for snapshotDate: Date) -> Date {
        MonthlyCheckInStore.effectiveMonth(for: snapshotDate, relativeTo: snapshotDate)
    }

    // MARK: - Chart Calculations

    private func calculateEvolutionData(from snapshots: [Snapshot]) {
        let sortedSnapshots = snapshots.sorted { $0.date < $1.date }
        let groupedByMonth = Dictionary(grouping: sortedSnapshots) { snapshot -> DateComponents in
            chartMonth(for: snapshot.date)
        }

        var series: [(date: Date, value: Decimal)] = []
        series.reserveCapacity(groupedByMonth.count)

        for (key, monthSnapshots) in groupedByMonth {
            var latestBySource: [UUID: Snapshot] = [:]
            for snapshot in monthSnapshots {
                guard let sourceId = snapshot.source?.id else { continue }
                if let existing = latestBySource[sourceId] {
                    if snapshot.date > existing.date {
                        latestBySource[sourceId] = snapshot
                    }
                } else {
                    latestBySource[sourceId] = snapshot
                }
            }

            let total = latestBySource.values.reduce(Decimal.zero) { $0 + $1.decimalValue }
            let date = Calendar.current.date(from: key) ?? Date()
            series.append((date: date, value: total))
        }

        series.sort { $0.date < $1.date }

        evolutionData = downsampleSeries(series, maxPoints: maxChartPoints)
    }

    private func calculateCategoryEvolutionData(from snapshots: [Snapshot], categories: [Category]) {
        let sortedSnapshots = snapshots.sorted { $0.date < $1.date }
        let categoriesWithData = Set(sortedSnapshots.compactMap { $0.source?.category?.id })
        let filteredCategories = categories.filter { categoriesWithData.contains($0.id) }
        let snapshotsByDay = Dictionary(grouping: sortedSnapshots) {
            Calendar.current.startOfDay(for: $0.date)
        }
        let uniqueDates = downsampleDates(snapshotsByDay.keys.sorted(), maxPoints: maxChartPoints)

        var latestBySource: [UUID: Snapshot] = [:]
        var points: [CategoryEvolutionPoint] = []

        for date in uniqueDates {
            if let daySnapshots = snapshotsByDay[date] {
                for snapshot in daySnapshots {
                    guard let sourceId = snapshot.source?.id else { continue }
                    latestBySource[sourceId] = snapshot
                }
            }

            var valuesByCategory: [UUID: Decimal] = [:]
            for snapshot in latestBySource.values {
                guard let category = snapshot.source?.category else { continue }
                valuesByCategory[category.id, default: 0] += snapshot.decimalValue
            }

            for category in filteredCategories {
                let value = valuesByCategory[category.id] ?? 0
                points.append(CategoryEvolutionPoint(
                    date: date,
                    categoryName: category.name,
                    colorHex: category.colorHex,
                    value: value
                ))
            }
        }

        categoryEvolutionData = points
    }

    private func calculateAllocationData(for sources: [InvestmentSource], breakdown: BreakdownMode) {
        switch breakdown {
        case .category:
            let categories = categoryRepository.categories
            let valuesByCategory = Dictionary(grouping: sources) { $0.category?.id ?? UUID() }

            allocationData = categories.compactMap { category in
                let categorySources = valuesByCategory[category.id] ?? []
                let categoryValue = categorySources.reduce(Decimal.zero) { $0 + $1.latestValue }
                guard categoryValue > 0 else { return nil }

                return (
                    category: category.name,
                    value: categoryValue,
                    color: category.colorHex
                )
            }.sorted { $0.value > $1.value }
        case .source:
            let colorMap = sourceColorMap(for: sources)
            allocationData = sources.compactMap { source in
                guard source.latestValue > 0 else { return nil }
                return (
                    category: source.name,
                    value: source.latestValue,
                    color: colorMap[source.id] ?? "#6B7280"
                )
            }.sorted { $0.value > $1.value }
        }
    }

    private func calculateAllocationEvolutionData(from snapshots: [Snapshot]) {
        let sortedSnapshots = snapshots.sorted { $0.date < $1.date }
        let groupedByMonth = Dictionary(grouping: sortedSnapshots) { snapshot -> DateComponents in
            chartMonth(for: snapshot.date)
        }

        let sortedMonths = groupedByMonth.keys.sorted {
            let d1 = ($0.year ?? 0) * 100 + ($0.month ?? 0)
            let d2 = ($1.year ?? 0) * 100 + ($1.month ?? 0)
            return d1 < d2
        }

        // First pass: compute totals per category across all months for stable ordering
        var globalCategoryTotals: [String: (total: Decimal, color: String)] = [:]
        var monthlyData: [(date: Date, categories: [String: (value: Decimal, color: String)])] = []

        for monthKey in sortedMonths {
            guard let monthSnapshots = groupedByMonth[monthKey],
                  let monthDate = Calendar.current.date(from: monthKey) else { continue }

            var categoryTotals: [String: (value: Decimal, color: String)] = [:]
            var sourceLatest: [UUID: Snapshot] = [:]

            for snapshot in monthSnapshots {
                guard let sourceId = snapshot.source?.id else { continue }
                if let existing = sourceLatest[sourceId] {
                    if snapshot.date > existing.date {
                        sourceLatest[sourceId] = snapshot
                    }
                } else {
                    sourceLatest[sourceId] = snapshot
                }
            }

            for (_, snapshot) in sourceLatest {
                let categoryName = snapshot.source?.category?.name ?? "Other"
                let colorHex = snapshot.source?.category?.colorHex ?? "#6B7280"
                let existing = categoryTotals[categoryName] ?? (value: 0, color: colorHex)
                categoryTotals[categoryName] = (value: existing.value + snapshot.decimalValue, color: colorHex)

                let globalExisting = globalCategoryTotals[categoryName] ?? (total: 0, color: colorHex)
                globalCategoryTotals[categoryName] = (total: globalExisting.total + snapshot.decimalValue, color: colorHex)
            }

            let total = categoryTotals.values.reduce(Decimal.zero) { $0 + $1.value }
            guard total > 0 else { continue }
            monthlyData.append((date: monthDate, categories: categoryTotals))
        }

        // Stable category order based on overall totals (largest first)
        let stableCategoryOrder = globalCategoryTotals
            .sorted { $0.value.total > $1.value.total }
            .map { $0.key }

        // Second pass: emit data points in stable order
        var result: [(date: Date, category: String, percentage: Double, color: String)] = []
        for month in monthlyData {
            let total = month.categories.values.reduce(Decimal.zero) { $0 + $1.value }
            guard total > 0 else { continue }

            for category in stableCategoryOrder {
                guard let info = month.categories[category] else { continue }
                let percentage = NSDecimalNumber(decimal: info.value / total * 100).doubleValue
                result.append((date: month.date, category: category, percentage: percentage, color: info.color))
            }
        }

        allocationEvolutionData = result
    }

    private func completedMonthKeys(
        sources: [InvestmentSource],
        snapshots: [Snapshot],
        after cutoff: Date
    ) -> Set<DateComponents> {
        let sourceIds = Set(sources.compactMap { $0.id })
        guard !sourceIds.isEmpty else { return [] }

        let groupedByMonth = Dictionary(grouping: snapshots) { snapshot -> DateComponents in
            chartMonth(for: snapshot.date)
        }

        var completed: Set<DateComponents> = []
        completed.reserveCapacity(groupedByMonth.count)

        for (key, monthSnapshots) in groupedByMonth {
            guard let monthDate = Calendar.current.date(from: key) else { continue }
            guard monthDate > cutoff else { continue }
            guard MonthlyCheckInStore.completionDate(for: monthDate) != nil else { continue }
            let monthSourceIds = Set(monthSnapshots.compactMap { $0.source?.id })
            if sourceIds.isSubset(of: monthSourceIds) {
                completed.insert(key)
            }
        }

        return completed
    }

    private func filterSnapshotsForCharts(
        sources: [InvestmentSource],
        snapshots: [Snapshot]
    ) -> [Snapshot] {
        guard let lastCompleted = MonthlyCheckInStore.latestCompletionDate()?.startOfMonth else {
            return snapshots
        }

        let completedMonthsAfter = completedMonthKeys(
            sources: sources,
            snapshots: snapshots,
            after: lastCompleted
        )

        return snapshots.filter { snapshot in
            let monthDate = chartMonthStart(for: snapshot.date)
            if monthDate <= lastCompleted {
                return true
            }
            let key = chartMonth(for: snapshot.date)
            return completedMonthsAfter.contains(key)
        }
    }

    private func groupSnapshotsBySource(_ snapshots: [Snapshot]) -> [UUID: [Snapshot]] {
        var grouped: [UUID: [Snapshot]] = [:]
        for snapshot in snapshots {
            guard let id = snapshot.source?.id else { continue }
            grouped[id, default: []].append(snapshot)
        }
        return grouped
    }

    private func calculatePerformanceData(
        for sources: [InvestmentSource],
        snapshotsBySource: [UUID: [Snapshot]],
        breakdown: BreakdownMode
    ) {
        switch breakdown {
        case .category:
            let categories = categoryRepository.categories
            let sourcesByCategory = Dictionary(grouping: sources) { $0.category?.id ?? UUID() }

            performanceData = categories.compactMap { category in
                let categorySources = sourcesByCategory[category.id] ?? []
                let snapshots = categorySources.compactMap { source -> [Snapshot]? in
                    let id = source.id
                    return snapshotsBySource[id]
                }.flatMap { $0 }
                guard snapshots.count >= 2 else { return nil }

                let monthlyTotals = monthlyTotalsByMonthYear(from: snapshots)
                guard let first = monthlyTotals.first,
                      let last = monthlyTotals.last,
                      first.totalValue > 0 else { return nil }
                let cagr = calculationService.calculateCAGR(
                    startValue: first.totalValue,
                    endValue: last.totalValue,
                    startDate: first.date,
                    endDate: last.date
                )

                return (
                    category: category.name,
                    cagr: cagr,
                    color: category.colorHex
                )
            }.sorted { $0.cagr > $1.cagr }
        case .source:
            let colorMap = sourceColorMap(for: sources)
            performanceData = sources.compactMap { source in
                guard let snapshots = snapshotsBySource[source.id], snapshots.count >= 2 else { return nil }
                let monthlyTotals = monthlyTotalsByMonthYear(from: snapshots)
                guard let first = monthlyTotals.first,
                      let last = monthlyTotals.last,
                      first.totalValue > 0 else { return nil }

                let cagr = calculationService.calculateCAGR(
                    startValue: first.totalValue,
                    endValue: last.totalValue,
                    startDate: first.date,
                    endDate: last.date
                )

                return (
                    category: source.name,
                    cagr: cagr,
                    color: colorMap[source.id] ?? "#6B7280"
                )
            }.sorted { $0.cagr > $1.cagr }
        }
    }

    private func calculateContributionsData(from snapshots: [Snapshot]) {
        let grouped = Dictionary(grouping: snapshots) { self.chartMonthStart(for: $0.date) }
        contributionsData = grouped.map { date, items in
            let total = items.reduce(Decimal.zero) { $0 + $1.decimalContribution }
            return (date: date, amount: total)
        }
        .sorted { $0.date < $1.date }
    }

    private func calculateRollingReturnData(from snapshots: [Snapshot]) {
        let monthlyTotals = monthlyTotals(from: snapshots)
        guard monthlyTotals.count >= 13 else {
            rollingReturnData = []
            return
        }

        var returns: [(date: Date, value: Double)] = []
        for index in 12..<monthlyTotals.count {
            let current = monthlyTotals[index]
            let base = monthlyTotals[index - 12]
            guard base.totalValue > 0 else { continue }
            let change = current.totalValue - base.totalValue
            let percent = NSDecimalNumber(decimal: change / base.totalValue).doubleValue * 100
            returns.append((date: current.date, value: percent))
        }
        rollingReturnData = returns
    }

    private func calculateRiskReturnData(
        for sources: [InvestmentSource],
        snapshotsBySource: [UUID: [Snapshot]]
    ) {
        let categories = categoryRepository.categories
        let sourcesByCategory = Dictionary(grouping: sources) { $0.category?.id ?? UUID() }

        riskReturnData = categories.compactMap { category in
            let categorySources = sourcesByCategory[category.id] ?? []
            let snapshots = categorySources.compactMap { source -> [Snapshot]? in
                let id = source.id
                return snapshotsBySource[id]
            }.flatMap { $0 }
            let monthlyTotals = monthlyTotalsByMonthYear(from: snapshots)
            guard monthlyTotals.count >= 3,
                  let first = monthlyTotals.first,
                  let last = monthlyTotals.last,
                  first.totalValue > 0 else { return nil }
            let cagr = calculationService.calculateCAGR(
                startValue: first.totalValue,
                endValue: last.totalValue,
                startDate: first.date,
                endDate: last.date
            )
            let monthlyReturns = monthlyReturnSeries(from: monthlyTotals)
            let volatility = calculationService.calculateVolatility(monthlyReturns: monthlyReturns)
            return (
                category: category.name,
                cagr: cagr,
                volatility: volatility,
                color: category.colorHex
            )
        }.sorted { $0.cagr > $1.cagr }
    }

    private func calculateCashflowData(from snapshots: [Snapshot]) {
        let monthlyTotals = monthlyTotals(from: snapshots)
        let contributionsByMonth = Dictionary(grouping: snapshots) { self.chartMonthStart(for: $0.date) }
            .mapValues { items in
                items.reduce(Decimal.zero) { $0 + $1.decimalContribution }
            }

        var data: [(date: Date, contributions: Decimal, netPerformance: Decimal)] = []
        for index in 0..<monthlyTotals.count {
            let current = monthlyTotals[index]
            let previousTotal = index > 0 ? monthlyTotals[index - 1].totalValue : 0
            let contributions = contributionsByMonth[current.date] ?? 0
            let netPerformance = current.totalValue - previousTotal - contributions
            data.append((date: current.date, contributions: contributions, netPerformance: netPerformance))
        }
        cashflowData = data
    }

    private func calculateDrawdownData(from snapshots: [Snapshot]) {
        let sortedSnapshots = snapshots.sorted { $0.date < $1.date }
        guard !sortedSnapshots.isEmpty else {
            drawdownData = []
            return
        }

        var peak = sortedSnapshots.first!.decimalValue
        var data: [(date: Date, drawdown: Double)] = []

        for snapshot in sortedSnapshots {
            let value = snapshot.decimalValue
            if value > peak {
                peak = value
            }

            let drawdown = peak > 0
                ? NSDecimalNumber(decimal: (peak - value) / peak).doubleValue * 100
                : 0

            data.append((date: snapshot.date, drawdown: -drawdown))
        }

        drawdownData = data
    }

    private func calculateVolatilityData(from snapshots: [Snapshot]) {
        let sortedSnapshots = snapshots.sorted { $0.date < $1.date }
        guard sortedSnapshots.count >= 3 else {
            volatilityData = []
            return
        }

        var data: [(date: Date, volatility: Double)] = []
        let windowSize = 3

        for i in windowSize..<sortedSnapshots.count {
            let window = Array(sortedSnapshots[(i - windowSize)..<i])
            let values = window.map { NSDecimalNumber(decimal: $0.decimalValue).doubleValue }

            let mean = values.reduce(0, +) / Double(values.count)
            let variance = values.map { pow($0 - mean, 2) }.reduce(0, +) / Double(values.count)
            let stdDev = sqrt(variance)
            let volatility = mean > 0 ? (stdDev / mean) * 100 : 0

            data.append((date: sortedSnapshots[i].date, volatility: volatility))
        }

        volatilityData = data
    }

    private func calculatePredictionData(from snapshots: [Snapshot]) {
        guard freemiumValidator.canViewPredictions() else {
            predictionData = []
            return
        }

        let result = predictionEngine.predict(snapshots: snapshots, monthsAhead: predictionMonthsAhead)
        predictionData = result.predictions
    }

    private func monthlyTotals(from snapshots: [Snapshot]) -> [(date: Date, totalValue: Decimal)] {
        let groupedByMonth = Dictionary(grouping: snapshots) { snapshot -> Date in
            chartMonthStart(for: snapshot.date)
        }

        let months = groupedByMonth.keys.sorted()
        guard !months.isEmpty else { return [] }

        var latestBySource: [UUID: Snapshot] = [:]
        var totals: [(date: Date, totalValue: Decimal)] = []

        for month in months {
            if let monthSnapshots = groupedByMonth[month] {
                for snapshot in monthSnapshots.sorted(by: { $0.date < $1.date }) {
                    guard let sourceId = snapshot.source?.id else { continue }
                    latestBySource[sourceId] = snapshot
                }
            }

            let total = latestBySource.values.reduce(Decimal.zero) { $0 + $1.decimalValue }
            totals.append((date: month, totalValue: total))
        }

        return totals
    }

    private func monthlyTotalsByMonthYear(from snapshots: [Snapshot]) -> [(date: Date, totalValue: Decimal)] {
        let sortedSnapshots = snapshots.sorted { $0.date < $1.date }
        let groupedByMonth = Dictionary(grouping: sortedSnapshots) { snapshot -> DateComponents in
            chartMonth(for: snapshot.date)
        }

        var totals: [(date: Date, totalValue: Decimal)] = []
        totals.reserveCapacity(groupedByMonth.count)

        for (key, monthSnapshots) in groupedByMonth {
            var latestBySource: [UUID: Snapshot] = [:]
            for snapshot in monthSnapshots {
                guard let sourceId = snapshot.source?.id else { continue }
                if let existing = latestBySource[sourceId] {
                    if snapshot.date > existing.date {
                        latestBySource[sourceId] = snapshot
                    }
                } else {
                    latestBySource[sourceId] = snapshot
                }
            }

            let total = latestBySource.values.reduce(Decimal.zero) { $0 + $1.decimalValue }
            let date = Calendar.current.date(from: key) ?? Date()
            totals.append((date: date, totalValue: total))
        }

        return totals.sorted { $0.date < $1.date }
    }

    private func monthlyReturnSeries(
        from monthlyTotals: [(date: Date, totalValue: Decimal)]
    ) -> [InvestmentMetrics.MonthlyReturn] {
        guard monthlyTotals.count >= 2 else { return [] }

        var returns: [InvestmentMetrics.MonthlyReturn] = []
        returns.reserveCapacity(monthlyTotals.count - 1)

        for index in 1..<monthlyTotals.count {
            let previous = monthlyTotals[index - 1]
            let current = monthlyTotals[index]
            guard previous.totalValue > 0 else { continue }
            let returnPercentage = NSDecimalNumber(
                decimal: (current.totalValue - previous.totalValue) / previous.totalValue
            ).doubleValue * 100
            returns.append(InvestmentMetrics.MonthlyReturn(
                date: current.date,
                returnPercentage: returnPercentage
            ))
        }

        return returns
    }

    func updatePredictionTargetDate(_ goals: [Goal]) {
        let futureGoalDates = goals.compactMap { $0.targetDate }.filter { $0 > Date() }
        guard let latestGoalDate = futureGoalDates.max(),
              let lastSnapshotDate = evolutionData.last?.date else {
            predictionMonthsAhead = 12
            return
        }

        let months = max(1, lastSnapshotDate.startOfMonth.monthsBetween(latestGoalDate.startOfMonth))
        predictionMonthsAhead = max(12, months)
        if selectedChartType == .prediction {
            updateChartData(chartType: selectedChartType, category: selectedCategory, timeRange: selectedTimeRange)
        }
    }

    // MARK: - Computed Properties

    var categories: [Category] {
        categoryRepository.categories
    }

    var availableChartTypes: [ChartType] {
        ChartType.allCases
    }

    var isPremium: Bool {
        freemiumValidator.isPremium
    }

    var hasData: Bool {
        !sourceRepository.sources.filter { shouldIncludeSource($0) }.isEmpty
    }
}