InvestmentTrackerApp/PortfolioJournal/Utilities/MonthlyCheckInStore.swift

import Foundation
import CoreData

// MonthlyCheckInStore — primary storage is CoreData (syncs via iCloud).
// UserDefaults is kept for one-time migration from older builds.
enum MonthlyCheckInStore {
    // Legacy UserDefaults keys (read-only after migration)
    private static let notesKey = "monthlyCheckInNotes"
    private static let completionsKey = "monthlyCheckInCompletions"
    private static let legacyLastCheckInKey = "lastCheckInDate"
    private static let entriesKey = "monthlyCheckInEntries"
    private static let migrationDoneKey = "journalMigratedToCoreData"
    static let graceDays = 20

    // MARK: - CoreData Context

    private static var context: NSManagedObjectContext {
        CoreDataStack.shared.viewContext
    }

    // MARK: - Public Accessors

    static func note(for date: Date) -> String {
        fetchEntry(for: monthKey(for: date))?.note ?? ""
    }

    static func setNote(_ note: String, for date: Date) {
        updateEntry(for: date) { entry in
            let trimmed = note.trimmingCharacters(in: .whitespacesAndNewlines)
            entry.note = trimmed.isEmpty ? nil : note
        }
    }

    static func rating(for date: Date) -> Int? {
        fetchEntry(for: monthKey(for: date))?.ratingValue
    }

    static func setRating(_ rating: Int?, for date: Date) {
        updateEntry(for: date) { entry in
            entry.ratingValue = rating
        }
    }

    static func mood(for date: Date) -> MonthlyCheckInMood? {
        fetchEntry(for: monthKey(for: date))?.mood
    }

    static func setMood(_ mood: MonthlyCheckInMood?, for date: Date) {
        updateEntry(for: date) { entry in
            entry.mood = mood
        }
    }

    static func monthKey(for date: Date) -> String {
        monthFormatter.string(from: effectiveMonth(for: date))
    }

    static func allNotes() -> [(date: Date, note: String)] {
        fetchAllEntries()
            .compactMap { entry in
                guard let key = entry.monthKey,
                      let date = monthFormatter.date(from: key) else { return nil }
                return (date: date, note: entry.note ?? "")
            }
            .sorted { $0.date > $1.date }
    }

    static func entry(for date: Date) -> MonthlyCheckInEntry? {
        fetchEntry(for: monthKey(for: date)).map(makeCheckInEntry)
    }

    static func allEntries() -> [(date: Date, entry: MonthlyCheckInEntry)] {
        fetchAllEntries()
            .compactMap { entry in
                guard let key = entry.monthKey,
                      let date = monthFormatter.date(from: key) else { return nil }
                return (date: date, entry: makeCheckInEntry(entry))
            }
            .sorted { $0.date > $1.date }
    }

    static func completionDate(for date: Date) -> Date? {
        fetchEntry(for: monthKey(for: date))?.completionTime
    }

    static func setCompletionDate(_ completionDate: Date, for month: Date) {
        let targetMonth = effectiveMonth(for: month, relativeTo: completionDate, graceDays: graceDays)
        let targetKey = monthFormatter.string(from: targetMonth)

        let calendar = Calendar.current
        if calendar.isDate(month, inSameDayAs: completionDate),
           calendar.component(.day, from: completionDate) > graceDays {

            let allExisting = fetchAllEntries()
            let completedPrevious = allExisting.compactMap { entry -> (month: Date, entry: JournalEntry)? in
                guard let key = entry.monthKey,
                      let entryMonth = monthFormatter.date(from: key)?.startOfMonth,
                      entry.completionTime != nil,
                      entryMonth < targetMonth else { return nil }
                return (month: entryMonth, entry: entry)
            }

            if let lastCompleted = completedPrevious.max(by: { $0.month < $1.month }) {
                var cursor = lastCompleted.month.adding(months: 1).startOfMonth
                while cursor < targetMonth {
                    let key = monthFormatter.string(from: cursor)
                    if fetchEntry(for: key) == nil {
                        let fallbackDate = min(cursor.endOfMonth, completionDate)
                        let new = JournalEntry(context: context)
                        new.monthKey = key
                        new.note = lastCompleted.entry.note
                        new.ratingValue = lastCompleted.entry.ratingValue
                        new.mood = lastCompleted.entry.mood
                        new.completionTime = fallbackDate
                    }
                    cursor = cursor.adding(months: 1).startOfMonth
                }
            }
        }

        // Ensure target month entry exists with the completion date.
        let targetEntry = fetchEntry(for: targetKey) ?? {
            let e = JournalEntry(context: context)
            e.monthKey = targetKey
            return e
        }()
        targetEntry.completionTime = completionDate

        // Backfill any previous entries that have no completion date.
        for entry in fetchAllEntries() {
            guard let key = entry.monthKey,
                  let entryMonth = monthFormatter.date(from: key)?.startOfMonth,
                  entryMonth < targetMonth,
                  entry.completionTime == nil else { continue }
            entry.completionTime = min(entryMonth.endOfMonth, completionDate)
        }

        saveContext()
    }

    static func latestCompletionDate() -> Date? {
        let request = JournalEntry.fetchRequest()
        request.predicate = NSPredicate(format: "completionTime != nil")
        request.sortDescriptors = [NSSortDescriptor(keyPath: \JournalEntry.completionTime, ascending: false)]
        request.fetchLimit = 1
        return (try? context.fetch(request))?.first?.completionTime
    }

    static func effectiveMonth(
        for date: Date,
        relativeTo referenceDate: Date = Date(),
        graceDays: Int = 20
    ) -> Date {
        let calendar = Calendar.current
        if calendar.isDate(date, inSameDayAs: referenceDate) {
            let day = calendar.component(.day, from: referenceDate)
            if day <= graceDays {
                return referenceDate.adding(months: -1).startOfMonth
            }
        }
        return date.startOfMonth
    }

    static func stats(referenceDate: Date = Date()) -> MonthlyCheckInStats {
        let cutoff = referenceDate.endOfMonth
        let entries = allEntries().filter { $0.date <= cutoff }
        let completions: [(month: Date, completion: Date, mood: MonthlyCheckInMood?)] = entries.compactMap { entry in
            guard let completion = entry.entry.completionDate else { return nil }
            return (month: entry.date.startOfMonth, completion: completion, mood: entry.entry.mood)
        }

        guard !completions.isEmpty else { return .empty }

        let deadlineDiffs = completions.map { item -> Double in
            let deadline = item.month.endOfMonth
            return deadline.timeIntervalSince(item.completion) / 86_400
        }

        let onTimeCompletions = completions.filter { item in
            item.completion <= item.month.endOfMonth
        }
        let onTimeMonths = Set(onTimeCompletions.map { $0.month })
        let totalCheckIns = completions.count
        let onTimeCount = onTimeMonths.count

        // The current month's check-in is still within its deadline (end of the
        // month), so having it pending must NOT read as a broken streak — that
        // made the card show 0x for most of every month. Start counting at the
        // current month when it is already done, otherwise at the previous one.
        // Any earlier month that is missing has passed its deadline and does
        // break the streak, so the walk backwards stays honest.
        var currentStreak = 0
        var cursor = referenceDate.startOfMonth
        if !onTimeMonths.contains(cursor) {
            cursor = cursor.adding(months: -1).startOfMonth
        }
        while onTimeMonths.contains(cursor) {
            currentStreak += 1
            cursor = cursor.adding(months: -1).startOfMonth
        }

        let sortedMonths = onTimeMonths.sorted()
        var bestStreak = 0
        var running = 0
        var previousMonth: Date?
        for month in sortedMonths {
            if let previousMonth, month == previousMonth.adding(months: 1).startOfMonth {
                running += 1
            } else {
                running = 1
            }
            bestStreak = max(bestStreak, running)
            previousMonth = month
        }

        let averageDaysBeforeDeadline = onTimeCount > 0
            ? deadlineDiffs.filter { $0 >= 0 }.average()
            : nil
        let closestCutoffDays = onTimeCount > 0
            ? deadlineDiffs.filter { $0 >= 0 }.min()
            : nil

        let recentMood = completions.sorted { $0.month > $1.month }.first?.mood
        let achievements = buildAchievements(
            currentStreak: currentStreak,
            bestStreak: bestStreak,
            onTimeCount: onTimeCount,
            totalCheckIns: totalCheckIns,
            closestCutoffDays: closestCutoffDays,
            averageDaysBeforeDeadline: averageDaysBeforeDeadline
        )

        return MonthlyCheckInStats(
            currentStreak: currentStreak,
            bestStreak: bestStreak,
            onTimeCount: onTimeCount,
            totalCheckIns: totalCheckIns,
            averageDaysBeforeDeadline: averageDaysBeforeDeadline,
            closestCutoffDays: closestCutoffDays,
            recentMood: recentMood,
            achievements: achievements
        )
    }

    static func achievementStatuses(referenceDate: Date = Date()) -> [MonthlyCheckInAchievementStatus] {
        let stats = stats(referenceDate: referenceDate)
        return achievementStatuses(for: stats)
    }

    static func clearAll() {
        for entry in fetchAllEntries() {
            context.delete(entry)
        }
        saveContext()
        let defaults = UserDefaults.standard
        defaults.removeObject(forKey: notesKey)
        defaults.removeObject(forKey: completionsKey)
        defaults.removeObject(forKey: entriesKey)
        defaults.removeObject(forKey: legacyLastCheckInKey)
    }

    // MARK: - One-time Migration from UserDefaults

    static func migrateIfNeeded() {
        guard !UserDefaults.standard.bool(forKey: migrationDoneKey) else { return }

        let legacyEntries = loadLegacyEntries()
        for (key, legacy) in legacyEntries {
            guard fetchEntry(for: key) == nil else { continue }
            let entry = JournalEntry(context: context)
            entry.monthKey = key
            entry.note = legacy.note
            entry.ratingValue = legacy.rating
            entry.mood = legacy.mood
            if let t = legacy.completionTime {
                entry.completionTime = Date(timeIntervalSince1970: t)
            }
            entry.createdAt = Date(timeIntervalSince1970: legacy.createdAt)
        }

        saveContext()
        UserDefaults.standard.set(true, forKey: migrationDoneKey)
    }

    // MARK: - Private CoreData Helpers

    private static func fetchEntry(for key: String) -> JournalEntry? {
        let request = JournalEntry.fetchRequest()
        request.predicate = NSPredicate(format: "monthKey == %@", key)
        request.fetchLimit = 1
        return try? context.fetch(request).first
    }

    private static func fetchAllEntries() -> [JournalEntry] {
        let request = JournalEntry.fetchRequest()
        return (try? context.fetch(request)) ?? []
    }

    private static func updateEntry(for date: Date, mutate: (JournalEntry) -> Void) {
        let key = monthKey(for: date)
        let entry = fetchEntry(for: key) ?? {
            let e = JournalEntry(context: context)
            e.monthKey = key
            return e
        }()

        mutate(entry)

        if entry.note?.trimmingCharacters(in: .whitespacesAndNewlines).isEmpty == true {
            entry.note = nil
        }

        let isEmpty = entry.note == nil
            && entry.ratingValue == nil
            && entry.mood == nil
            && entry.completionTime == nil
        if isEmpty {
            context.delete(entry)
        }

        saveContext()
    }

    private static func saveContext() {
        guard context.hasChanges else { return }
        try? context.save()
    }

    private static func makeCheckInEntry(_ entry: JournalEntry) -> MonthlyCheckInEntry {
        MonthlyCheckInEntry(
            note: entry.note,
            rating: entry.ratingValue,
            mood: entry.mood,
            completionTime: entry.completionTime?.timeIntervalSince1970,
            createdAt: entry.createdAt?.timeIntervalSince1970 ?? Date().timeIntervalSince1970
        )
    }

    // MARK: - Legacy UserDefaults Reader (for migration only)

    private static func loadLegacyEntries() -> [String: MonthlyCheckInEntry] {
        if let data = UserDefaults.standard.data(forKey: entriesKey),
           let decoded = try? JSONDecoder().decode([String: MonthlyCheckInEntry].self, from: data),
           !decoded.isEmpty {
            return decoded
        }
        return loadAndMergeLegacyKeys()
    }

    private static func loadAndMergeLegacyKeys() -> [String: MonthlyCheckInEntry] {
        let notes = loadLegacyNotes()
        let completions = loadLegacyCompletions()
        guard !notes.isEmpty || !completions.isEmpty else { return [:] }

        var entries: [String: MonthlyCheckInEntry] = [:]
        let now = Date().timeIntervalSince1970
        for (key, note) in notes {
            entries[key] = MonthlyCheckInEntry(
                note: note, rating: nil, mood: nil,
                completionTime: completions[key], createdAt: now
            )
        }
        for (key, completion) in completions where entries[key] == nil {
            entries[key] = MonthlyCheckInEntry(
                note: nil, rating: nil, mood: nil,
                completionTime: completion, createdAt: completion
            )
        }
        return entries
    }

    private static func loadLegacyNotes() -> [String: String] {
        guard let data = UserDefaults.standard.data(forKey: notesKey),
              let decoded = try? JSONDecoder().decode([String: String].self, from: data) else { return [:] }
        return decoded
    }

    private static func loadLegacyCompletions() -> [String: Double] {
        guard let data = UserDefaults.standard.data(forKey: completionsKey),
              let decoded = try? JSONDecoder().decode([String: Double].self, from: data) else { return [:] }
        return decoded
    }

    // MARK: - Private Achievement Helpers

    private struct MonthlyCheckInAchievementRule {
        let achievement: MonthlyCheckInAchievement
        let isUnlocked: (Int, Int, Int, Int, Double?, Double?) -> Bool
    }

    private static let achievementRules: [MonthlyCheckInAchievementRule] = [
        MonthlyCheckInAchievementRule(
            achievement: MonthlyCheckInAchievement(
                key: "streak_3",
                title: String(localized: "achievement_streak_3_title"),
                detail: String(localized: "achievement_streak_3_detail"),
                icon: "flame.fill"
            ),
            isUnlocked: { currentStreak, _, _, _, _, _ in currentStreak >= 3 }
        ),
        MonthlyCheckInAchievementRule(
            achievement: MonthlyCheckInAchievement(
                key: "streak_6",
                title: String(localized: "achievement_streak_6_title"),
                detail: String(localized: "achievement_streak_6_detail"),
                icon: "bolt.heart.fill"
            ),
            isUnlocked: { currentStreak, _, _, _, _, _ in currentStreak >= 6 }
        ),
        MonthlyCheckInAchievementRule(
            achievement: MonthlyCheckInAchievement(
                key: "streak_12",
                title: String(localized: "achievement_streak_12_title"),
                detail: String(localized: "achievement_streak_12_detail"),
                icon: "calendar.circle.fill"
            ),
            isUnlocked: { _, bestStreak, _, _, _, _ in bestStreak >= 12 }
        ),
        MonthlyCheckInAchievementRule(
            achievement: MonthlyCheckInAchievement(
                key: "perfect_on_time",
                title: String(localized: "achievement_perfect_on_time_title"),
                detail: String(localized: "achievement_perfect_on_time_detail"),
                icon: "checkmark.seal.fill"
            ),
            isUnlocked: { _, _, onTimeCount, totalCheckIns, _, _ in
                onTimeCount == totalCheckIns && totalCheckIns >= 3
            }
        ),
        MonthlyCheckInAchievementRule(
            achievement: MonthlyCheckInAchievement(
                key: "clutch_finish",
                title: String(localized: "achievement_clutch_finish_title"),
                detail: String(localized: "achievement_clutch_finish_detail"),
                icon: "hourglass"
            ),
            isUnlocked: { _, _, _, _, closestCutoffDays, _ in
                if let closestCutoffDays { return closestCutoffDays <= 2 }
                return false
            }
        ),
        MonthlyCheckInAchievementRule(
            achievement: MonthlyCheckInAchievement(
                key: "early_bird",
                title: String(localized: "achievement_early_bird_title"),
                detail: String(localized: "achievement_early_bird_detail"),
                icon: "sun.max.fill"
            ),
            isUnlocked: { _, _, _, totalCheckIns, _, averageDaysBeforeDeadline in
                if let averageDaysBeforeDeadline {
                    return averageDaysBeforeDeadline >= 10 && totalCheckIns >= 3
                }
                return false
            }
        )
    ]

    private static func achievementStatuses(for stats: MonthlyCheckInStats) -> [MonthlyCheckInAchievementStatus] {
        achievementRules.map { rule in
            MonthlyCheckInAchievementStatus(
                achievement: rule.achievement,
                isUnlocked: rule.isUnlocked(
                    stats.currentStreak,
                    stats.bestStreak,
                    stats.onTimeCount,
                    stats.totalCheckIns,
                    stats.closestCutoffDays,
                    stats.averageDaysBeforeDeadline
                )
            )
        }
    }

    private static func buildAchievements(
        currentStreak: Int,
        bestStreak: Int,
        onTimeCount: Int,
        totalCheckIns: Int,
        closestCutoffDays: Double?,
        averageDaysBeforeDeadline: Double?
    ) -> [MonthlyCheckInAchievement] {
        achievementRules.compactMap { rule in
            rule.isUnlocked(
                currentStreak, bestStreak, onTimeCount, totalCheckIns,
                closestCutoffDays, averageDaysBeforeDeadline
            ) ? rule.achievement : nil
        }
    }

    private static var monthFormatter: DateFormatter {
        let formatter = DateFormatter()
        formatter.dateFormat = "yyyy-MM"
        return formatter
    }
}