Reflect/Shared/Views/InsightsView/ReportsViewModel.swift

//
//  ReportsViewModel.swift
//  Reflect
//
//  ViewModel for AI mood report generation and PDF export.
//

import Foundation
import SwiftUI
import FoundationModels

// MARK: - Generation State

enum ReportGenerationState: Equatable {
    case idle
    case generating
    case completed
    case failed(String)
}

// MARK: - ViewModel

@MainActor
class ReportsViewModel: ObservableObject {

    // MARK: - Published State

    @Published var startDate: Date = Calendar.current.date(byAdding: .month, value: -1, to: Date()) ?? Date()
    @Published var endDate: Date = Date()
    @Published var reportType: ReportType = .quickSummary
    @Published var generationState: ReportGenerationState = .idle
    @Published var progressValue: Double = 0.0
    @Published var progressMessage: String = ""
    @Published var generatedReport: MoodReport?
    @Published var exportedPDFURL: URL?
    @Published var showShareSheet: Bool = false
    @Published var showPrivacyConfirmation: Bool = false
    @Published var errorMessage: String?
    @Published var isAIAvailable: Bool = false

    // MARK: - Computed Properties

    var entriesInRange: [MoodEntryModel] {
        let allEntries = DataController.shared.getData(
            startDate: Calendar.current.startOfDay(for: startDate),
            endDate: Calendar.current.startOfDay(for: Calendar.current.date(byAdding: .day, value: 1, to: endDate) ?? endDate),
            includedDays: [1, 2, 3, 4, 5, 6, 7]
        )
        return allEntries.filter { ![.missing, .placeholder].contains($0.mood) }
    }

    var validEntryCount: Int { entriesInRange.count }

    var canGenerate: Bool {
        validEntryCount >= 3 && isAIAvailable
    }

    var daySpan: Int {
        let components = Calendar.current.dateComponents([.day], from: startDate, to: endDate)
        return (components.day ?? 0) + 1
    }

    // MARK: - Dependencies

    private var insightService: Any?
    private let summarizer = MoodDataSummarizer()
    private let pdfGenerator = ReportPDFGenerator()
    private let calendar = Calendar.current
    private var generationTask: Task<Void, Never>?

    // MARK: - Initialization

    init() {
        if #available(iOS 26, *) {
            let service = FoundationModelsInsightService()
            insightService = service
            isAIAvailable = service.isAvailable
        } else {
            insightService = nil
            isAIAvailable = false
        }
    }

    deinit {
        generationTask?.cancel()
    }

    // MARK: - Report Generation

    func generateReport() {
        generationTask?.cancel()
        generationState = .generating
        progressValue = 0.0
        progressMessage = String(localized: "Preparing data...")
        errorMessage = nil
        generatedReport = nil

        generationTask = Task {
            do {
                let entries = entriesInRange
                let reportEntries = entries.map { ReportEntry(from: $0) }

                let report: MoodReport
                switch reportType {
                case .quickSummary:
                    report = try await generateQuickSummary(entries: entries, reportEntries: reportEntries)
                case .detailed:
                    report = try await generateDetailedReport(entries: entries, reportEntries: reportEntries)
                }

                guard !Task.isCancelled else { return }

                generatedReport = report
                generationState = .completed
                progressValue = 1.0
                progressMessage = String(localized: "Report ready")

                AnalyticsManager.shared.track(.reportGenerated(
                    type: reportType.rawValue,
                    entryCount: validEntryCount,
                    daySpan: daySpan
                ))
            } catch {
                guard !Task.isCancelled else { return }
                generationState = .failed(error.localizedDescription)
                errorMessage = error.localizedDescription
                AnalyticsManager.shared.track(.reportGenerationFailed(error: error.localizedDescription))
            }
        }
    }

    func cancelGeneration() {
        generationTask?.cancel()
        generationState = .idle
        progressValue = 0.0
        progressMessage = ""
        AnalyticsManager.shared.track(.reportCancelled)
    }

    // MARK: - PDF Export

    func exportPDF() {
        guard let report = generatedReport else { return }

        Task {
            do {
                let url = try await pdfGenerator.generatePDF(from: report)
                exportedPDFURL = url
                showShareSheet = true

                AnalyticsManager.shared.track(.reportExported(
                    type: reportType.rawValue,
                    entryCount: validEntryCount
                ))
            } catch {
                errorMessage = error.localizedDescription
                generationState = .failed(error.localizedDescription)
            }
        }
    }

    // MARK: - Quick Summary Generation

    private func generateQuickSummary(entries: [MoodEntryModel], reportEntries: [ReportEntry]) async throws -> MoodReport {
        let overview = buildOverview(entries: entries)
        let weeks = splitIntoWeeks(entries: reportEntries)

        progressValue = 0.3
        progressMessage = String(localized: "Generating AI summary...")

        var quickSummaryText: String?

        if #available(iOS 26, *) {
            let summary = summarizer.summarize(entries: entries, periodName: "selected period")
            let promptData = summarizer.toPromptString(summary)

            let session = LanguageModelSession(instructions: clinicalSystemInstructions)

            let prompt = """
            Analyze this mood data and generate a clinical summary report:

            \(promptData)

            Generate a factual, third-person clinical summary suitable for sharing with a therapist.
            """

            do {
                let response = try await session.respond(to: prompt, generating: AIQuickSummaryResponse.self)

                guard !Task.isCancelled else { throw CancellationError() }

                var text = response.content.summary
                if !response.content.keyObservations.isEmpty {
                    text += "\n\nKey Observations:\n" + response.content.keyObservations.map { "- \($0)" }.joined(separator: "\n")
                }
                if !response.content.recommendations.isEmpty {
                    text += "\n\nRecommendations:\n" + response.content.recommendations.map { "- \($0)" }.joined(separator: "\n")
                }
                quickSummaryText = text
            } catch is CancellationError {
                throw CancellationError()
            } catch {
                quickSummaryText = "Summary unavailable: \(error.localizedDescription)"
            }
        }

        progressValue = 0.9

        let monthlySummaries = buildMonthlySummaries(entries: reportEntries)
        let yearlySummaries = buildYearlySummaries(entries: reportEntries)

        return MoodReport(
            reportType: .quickSummary,
            generatedAt: Date(),
            overview: overview,
            weeks: weeks,
            monthlySummaries: monthlySummaries,
            yearlySummaries: yearlySummaries,
            quickSummary: quickSummaryText
        )
    }

    // MARK: - Detailed Report Generation

    private func generateDetailedReport(entries: [MoodEntryModel], reportEntries: [ReportEntry]) async throws -> MoodReport {
        let overview = buildOverview(entries: entries)
        var weeks = splitIntoWeeks(entries: reportEntries)
        var monthlySummaries = buildMonthlySummaries(entries: reportEntries)
        var yearlySummaries = buildYearlySummaries(entries: reportEntries)

        let totalSections = weeks.count + monthlySummaries.count + yearlySummaries.count
        var completedSections = 0

        // Generate weekly AI summaries — batched at 4 concurrent
        if #available(iOS 26, *) {
            let batchSize = 4

            for batchStart in stride(from: 0, to: weeks.count, by: batchSize) {
                guard !Task.isCancelled else { throw CancellationError() }

                let batchEnd = min(batchStart + batchSize, weeks.count)
                let batchIndices = batchStart..<batchEnd

                await withTaskGroup(of: (Int, String?).self) { group in
                    for index in batchIndices {
                        group.addTask { @MainActor in
                            let week = weeks[index]
                            let summary = await self.generateWeeklySummary(week: week)
                            return (index, summary)
                        }
                    }

                    for await (index, summary) in group {
                        weeks[index].aiSummary = summary
                        completedSections += 1
                        progressValue = Double(completedSections) / Double(totalSections)
                        progressMessage = String(localized: "Generating weekly summary \(completedSections) of \(weeks.count)...")
                    }
                }
            }

            // Generate monthly AI summaries — concurrent
            guard !Task.isCancelled else { throw CancellationError() }
            progressMessage = String(localized: "Generating monthly summaries...")

            await withTaskGroup(of: (Int, String?).self) { group in
                for (index, monthSummary) in monthlySummaries.enumerated() {
                    group.addTask { @MainActor in
                        let summary = await self.generateMonthlySummary(month: monthSummary, allEntries: reportEntries)
                        return (index, summary)
                    }
                }

                for await (index, summary) in group {
                    monthlySummaries[index].aiSummary = summary
                    completedSections += 1
                    progressValue = Double(completedSections) / Double(totalSections)
                }
            }

            // Generate yearly AI summaries — concurrent
            guard !Task.isCancelled else { throw CancellationError() }

            if !yearlySummaries.isEmpty {
                progressMessage = String(localized: "Generating yearly summaries...")

                await withTaskGroup(of: (Int, String?).self) { group in
                    for (index, yearSummary) in yearlySummaries.enumerated() {
                        group.addTask { @MainActor in
                            let summary = await self.generateYearlySummary(year: yearSummary, allEntries: reportEntries)
                            return (index, summary)
                        }
                    }

                    for await (index, summary) in group {
                        yearlySummaries[index].aiSummary = summary
                        completedSections += 1
                        progressValue = Double(completedSections) / Double(totalSections)
                    }
                }
            }
        }

        return MoodReport(
            reportType: .detailed,
            generatedAt: Date(),
            overview: overview,
            weeks: weeks,
            monthlySummaries: monthlySummaries,
            yearlySummaries: yearlySummaries,
            quickSummary: nil
        )
    }

    // MARK: - AI Summary Generators

    @available(iOS 26, *)
    private func generateWeeklySummary(week: ReportWeek) async -> String? {
        let session = LanguageModelSession(instructions: clinicalSystemInstructions)

        let moodList = week.entries.sorted(by: { $0.date < $1.date }).map { entry in
            let day = entry.date.formatted(.dateTime.weekday(.abbreviated))
            let mood = entry.mood.widgetDisplayName
            let notes = entry.notes ?? "no notes"
            return "\(day): \(mood) (\(notes))"
        }.joined(separator: "\n")

        let prompt = """
        Summarize this week's mood data (Week \(week.weekNumber)):
        \(moodList)

        Average mood: \(String(format: "%.1f", weekAverage(week)))/5
        """

        do {
            let response = try await session.respond(to: prompt, generating: AIWeeklySummary.self)
            return response.content.summary
        } catch {
            return "Summary unavailable"
        }
    }

    @available(iOS 26, *)
    private func generateMonthlySummary(month: ReportMonthSummary, allEntries: [ReportEntry]) async -> String? {
        let session = LanguageModelSession(instructions: clinicalSystemInstructions)

        let monthEntries = allEntries.filter {
            calendar.component(.month, from: $0.date) == month.month &&
            calendar.component(.year, from: $0.date) == month.year
        }

        let moodDist = Dictionary(grouping: monthEntries, by: { $0.mood.widgetDisplayName })
            .mapValues { $0.count }
            .sorted { $0.value > $1.value }
            .map { "\($0.key): \($0.value)" }
            .joined(separator: ", ")

        let prompt = """
        Summarize this month's mood data (\(month.title)):
        \(month.entryCount) entries, average mood: \(String(format: "%.1f", month.averageMood))/5
        Distribution: \(moodDist)
        """

        do {
            let response = try await session.respond(to: prompt, generating: AIMonthSummary.self)
            return response.content.summary
        } catch {
            return "Summary unavailable"
        }
    }

    @available(iOS 26, *)
    private func generateYearlySummary(year: ReportYearSummary, allEntries: [ReportEntry]) async -> String? {
        let session = LanguageModelSession(instructions: clinicalSystemInstructions)

        let yearEntries = allEntries.filter { calendar.component(.year, from: $0.date) == year.year }

        let monthlyAvgs = Dictionary(grouping: yearEntries) { calendar.component(.month, from: $0.date) }
            .sorted { $0.key < $1.key }
            .map { (month, entries) in
                let avg = Double(entries.reduce(0) { $0 + Int($1.mood.rawValue) + 1 }) / Double(entries.count)
                let formatter = DateFormatter()
                formatter.dateFormat = "MMM"
                var comps = DateComponents()
                comps.month = month
                let date = calendar.date(from: comps) ?? Date()
                return "\(formatter.string(from: date)): \(String(format: "%.1f", avg))"
            }
            .joined(separator: ", ")

        let prompt = """
        Summarize this year's mood data (\(year.year)):
        \(year.entryCount) entries, average mood: \(String(format: "%.1f", year.averageMood))/5
        Monthly averages: \(monthlyAvgs)
        """

        do {
            let response = try await session.respond(to: prompt, generating: AIYearSummary.self)
            return response.content.summary
        } catch {
            return "Summary unavailable"
        }
    }

    // MARK: - Data Building Helpers

    private func buildOverview(entries: [MoodEntryModel]) -> ReportOverviewStats {
        let validEntries = entries.filter { ![.missing, .placeholder].contains($0.mood) }
        let total = validEntries.count
        let avgMood = total > 0 ? Double(validEntries.reduce(0) { $0 + Int($1.moodValue) + 1 }) / Double(total) : 0

        var distribution: [Mood: Int] = [:]
        for entry in validEntries {
            distribution[entry.mood, default: 0] += 1
        }

        let sorted = validEntries.sorted { $0.forDate < $1.forDate }
        let trend: String
        if sorted.count >= 4 {
            let half = sorted.count / 2
            let firstAvg = Double(sorted.prefix(half).reduce(0) { $0 + Int($1.moodValue) + 1 }) / Double(half)
            let secondAvg = Double(sorted.suffix(half).reduce(0) { $0 + Int($1.moodValue) + 1 }) / Double(half)
            let diff = secondAvg - firstAvg
            trend = diff > 0.5 ? "Improving" : diff < -0.5 ? "Declining" : "Stable"
        } else {
            trend = "Stable"
        }

        let dateFormatter = DateFormatter()
        dateFormatter.dateStyle = .medium
        let rangeStr: String
        if let first = sorted.first, let last = sorted.last {
            rangeStr = "\(dateFormatter.string(from: first.forDate)) - \(dateFormatter.string(from: last.forDate))"
        } else {
            rangeStr = "No data"
        }

        return ReportOverviewStats(
            totalEntries: total,
            averageMood: avgMood,
            moodDistribution: distribution,
            trend: trend,
            dateRange: rangeStr
        )
    }

    private func splitIntoWeeks(entries: [ReportEntry]) -> [ReportWeek] {
        let sorted = entries.sorted { $0.date < $1.date }
        guard let firstDate = sorted.first?.date else { return [] }

        var weeks: [ReportWeek] = []
        var weekStart = calendar.startOfDay(for: firstDate)
        var weekNumber = 1

        while weekStart <= (sorted.last?.date ?? Date()) {
            let weekEnd = calendar.date(byAdding: .day, value: 6, to: weekStart) ?? weekStart
            let weekEntries = sorted.filter { entry in
                let entryDay = calendar.startOfDay(for: entry.date)
                return entryDay >= weekStart && entryDay <= weekEnd
            }

            if !weekEntries.isEmpty {
                weeks.append(ReportWeek(
                    weekNumber: weekNumber,
                    startDate: weekStart,
                    endDate: weekEnd,
                    entries: weekEntries
                ))
            }

            weekStart = calendar.date(byAdding: .day, value: 7, to: weekStart) ?? weekStart
            weekNumber += 1
        }

        return weeks
    }

    private func buildMonthlySummaries(entries: [ReportEntry]) -> [ReportMonthSummary] {
        let grouped = Dictionary(grouping: entries) { entry in
            let month = calendar.component(.month, from: entry.date)
            let year = calendar.component(.year, from: entry.date)
            return "\(year)-\(month)"
        }

        return grouped.map { (key, monthEntries) in
            let components = key.split(separator: "-")
            let year = Int(components[0]) ?? 0
            let month = Int(components[1]) ?? 0
            let avg = Double(monthEntries.reduce(0) { $0 + Int($1.mood.rawValue) + 1 }) / Double(monthEntries.count)

            return ReportMonthSummary(
                month: month,
                year: year,
                entryCount: monthEntries.count,
                averageMood: avg
            )
        }
        .sorted { ($0.year, $0.month) < ($1.year, $1.month) }
    }

    private func buildYearlySummaries(entries: [ReportEntry]) -> [ReportYearSummary] {
        let grouped = Dictionary(grouping: entries) { calendar.component(.year, from: $0.date) }
        guard grouped.count > 1 else { return [] } // Only generate if range spans multiple years

        return grouped.map { (year, yearEntries) in
            let avg = Double(yearEntries.reduce(0) { $0 + Int($1.mood.rawValue) + 1 }) / Double(yearEntries.count)
            return ReportYearSummary(year: year, entryCount: yearEntries.count, averageMood: avg)
        }
        .sorted { $0.year < $1.year }
    }

    private func weekAverage(_ week: ReportWeek) -> Double {
        let total = week.entries.reduce(0) { $0 + Int($1.mood.rawValue) + 1 }
        return week.entries.isEmpty ? 0 : Double(total) / Double(week.entries.count)
    }

    // MARK: - Clinical System Instructions

    private var clinicalSystemInstructions: String {
        let languageCode = Locale.current.language.languageCode?.identifier ?? "en"
        return """
        You are a clinical mood data analyst generating a professional mood report. \
        Use third-person perspective (e.g., "The individual", "The subject"). \
        Be factual, neutral, and objective. Do not use casual language, emojis, or personality-driven tone. \
        Reference specific data points and patterns. \
        This report may be shared with a therapist or healthcare professional. \
        Generate all text in the language with code: \(languageCode).
        """
    }
}