InvestmentTrackerApp/InvestmentTracker/ViewModels/ChartsViewModel.swift

import Foundation
import Combine

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

    enum ChartType: String, CaseIterable, Identifiable {
        case evolution = "Evolution"
        case allocation = "Allocation"
        case performance = "Performance"
        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 .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, .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 .drawdown:
                return "Analyze declines from peak values"
            case .volatility:
                return "Understand investment risk levels"
            case .prediction:
                return "View 12-month forecasts"
            }
        }
    }

    // MARK: - Published Properties

    @Published var selectedChartType: ChartType = .evolution
    @Published var selectedCategory: Category?
    @Published var selectedTimeRange: TimeRange = .year

    @Published var evolutionData: [(date: Date, value: Decimal)] = []
    @Published var allocationData: [(category: String, value: Decimal, color: String)] = []
    @Published var performanceData: [(category: String, cagr: Double, color: String)] = []
    @Published var drawdownData: [(date: Date, drawdown: Double)] = []
    @Published var volatilityData: [(date: Date, volatility: Double)] = []
    @Published var predictionData: [Prediction] = []

    @Published var isLoading = false
    @Published var showingPaywall = false

    // MARK: - Time Range

    enum TimeRange: String, CaseIterable, Identifiable {
        case month = "1M"
        case quarter = "3M"
        case halfYear = "6M"
        case year = "1Y"
        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 .all: return nil
            }
        }
    }

    // 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 var cancellables = Set<AnyCancellable>()

    // MARK: - Initialization

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

        setupObservers()
        loadData()
    }

    // MARK: - Setup

    private func setupObservers() {
        Publishers.CombineLatest3($selectedChartType, $selectedCategory, $selectedTimeRange)
            .debounce(for: .milliseconds(100), scheduler: DispatchQueue.main)
            .sink { [weak self] chartType, category, timeRange in
                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
        FirebaseService.shared.logChartViewed(
            chartType: chartType.rawValue,
            isPremium: chartType.isPremium
        )
    }

    private func updateChartData(chartType: ChartType, category: Category?, timeRange: TimeRange) {
        isLoading = true

        let sources: [InvestmentSource]
        if let category = category {
            sources = sourceRepository.fetchSources(for: category)
        } else {
            sources = sourceRepository.sources
        }

        var snapshots = sources.flatMap { $0.snapshotsArray }

        // Filter by time range
        if let months = timeRange.months {
            let cutoff = Calendar.current.date(byAdding: .month, value: -months, to: Date()) ?? Date()
            snapshots = snapshots.filter { $0.date >= cutoff }
        }

        // Apply freemium filter
        snapshots = freemiumValidator.filterSnapshots(snapshots)

        switch chartType {
        case .evolution:
            calculateEvolutionData(from: snapshots)
        case .allocation:
            calculateAllocationData()
        case .performance:
            calculatePerformanceData()
        case .drawdown:
            calculateDrawdownData(from: snapshots)
        case .volatility:
            calculateVolatilityData(from: snapshots)
        case .prediction:
            calculatePredictionData(from: snapshots)
        }

        isLoading = false
    }

    // MARK: - Chart Calculations

    private func calculateEvolutionData(from snapshots: [Snapshot]) {
        let sortedSnapshots = snapshots.sorted { $0.date < $1.date }

        var dateValues: [Date: Decimal] = [:]

        for snapshot in sortedSnapshots {
            let day = Calendar.current.startOfDay(for: snapshot.date)
            dateValues[day, default: 0] += snapshot.decimalValue
        }

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

    private func calculateAllocationData() {
        let categories = categoryRepository.categories
        let total = categories.reduce(Decimal.zero) { $0 + $1.totalValue }

        allocationData = categories
            .filter { $0.totalValue > 0 }
            .map { category in
                (
                    category: category.name,
                    value: category.totalValue,
                    color: category.colorHex
                )
            }
            .sorted { $0.value > $1.value }
    }

    private func calculatePerformanceData() {
        let categories = categoryRepository.categories

        performanceData = categories.compactMap { category in
            let snapshots = category.sourcesArray.flatMap { $0.snapshotsArray }
            guard snapshots.count >= 2 else { return nil }

            let metrics = calculationService.calculateMetrics(for: snapshots)

            return (
                category: category.name,
                cagr: metrics.cagr,
                color: category.colorHex
            )
        }.sorted { $0.cagr > $1.cagr }
    }

    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)
        predictionData = result.predictions
    }

    // MARK: - Computed Properties

    var categories: [Category] {
        categoryRepository.categories
    }

    var availableChartTypes: [ChartType] {
        ChartType.allCases
    }

    var isPremium: Bool {
        freemiumValidator.isPremium
    }

    var hasData: Bool {
        !sourceRepository.sources.isEmpty
    }
}