import Foundation

enum MonthlyCheckInStore {
    private static let notesKey = "monthlyCheckInNotes"
    private static let completionsKey = "monthlyCheckInCompletions"
    private static let legacyLastCheckInKey = "lastCheckInDate"
    private static let entriesKey = "monthlyCheckInEntries"
    private static let graceDays = 20

    // MARK: - Public Accessors

    static func note(for date: Date) -> String {
        entry(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? {
        entry(for: date)?.rating
    }

    static func setRating(_ rating: Int?, for date: Date) {
        updateEntry(for: date) { entry in
            if let rating, rating > 0 {
                entry.rating = min(max(1, rating), 5)
            } else {
                entry.rating = nil
            }
        }
    }

    static func mood(for date: Date) -> MonthlyCheckInMood? {
        entry(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 {
        Self.monthFormatter.string(from: effectiveMonth(for: date))
    }

    static func allNotes() -> [(date: Date, note: String)] {
        loadEntries()
            .compactMap { key, entry in
                guard let date = Self.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? {
        loadEntries()[monthKey(for: date)]
    }

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

    static func completionDate(for date: Date) -> Date? {
        entry(for: date)?.completionDate
    }

    static func setCompletionDate(_ completionDate: Date, for month: Date) {
        let targetMonth = effectiveMonth(for: month, relativeTo: completionDate, graceDays: graceDays)
        let targetKey = monthKey(for: targetMonth)
        var entries = loadEntries()
        var didChange = false

        let calendar = Calendar.current
        if calendar.isDate(month, inSameDayAs: completionDate),
           calendar.component(.day, from: completionDate) > graceDays {
            let completedEntries = entries.compactMap { key, entry -> (month: Date, entry: MonthlyCheckInEntry)? in
                guard let monthDate = monthFormatter.date(from: key)?.startOfMonth else { return nil }
                guard entry.completionTime != nil else { return nil }
                guard monthDate < targetMonth else { return nil }
                return (month: monthDate, entry: entry)
            }

            if let lastCompleted = completedEntries.max(by: { $0.month < $1.month }) {
                var cursor = lastCompleted.month.adding(months: 1).startOfMonth
                while cursor < targetMonth {
                    let key = monthFormatter.string(from: cursor)
                    if entries[key] == nil {
                        let fallbackCompletion = min(cursor.endOfMonth, completionDate)
                        let entry = MonthlyCheckInEntry(
                            note: lastCompleted.entry.note,
                            rating: lastCompleted.entry.rating,
                            mood: lastCompleted.entry.mood,
                            completionTime: fallbackCompletion.timeIntervalSince1970,
                            createdAt: Date().timeIntervalSince1970
                        )
                        entries[key] = entry
                        didChange = true
                    }
                    cursor = cursor.adding(months: 1).startOfMonth
                }
            }
        }

        // Ensure the target month entry exists.
        var targetEntry = entries[targetKey] ?? MonthlyCheckInEntry(
            note: nil,
            rating: nil,
            mood: nil,
            completionTime: legacyCompletion(for: targetKey),
            createdAt: Date().timeIntervalSince1970
        )
        targetEntry.completionTime = completionDate.timeIntervalSince1970
        entries[targetKey] = targetEntry
        didChange = true

        // Mark all previous pending entries as completed as well.
        for (key, entry) in entries {
            guard let entryMonth = monthFormatter.date(from: key)?.startOfMonth else { continue }
            guard entryMonth < targetMonth else { continue }
            guard entry.completionTime == nil else { continue }

            var updated = entry
            let fallbackCompletion = min(entryMonth.endOfMonth, completionDate)
            updated.completionTime = fallbackCompletion.timeIntervalSince1970
            entries[key] = updated
            didChange = true
        }

        guard didChange else { return }
        saveEntries(entries)
        persistLegacyMirrors(entries)
    }

    static func latestCompletionDate() -> Date? {
        let latestEntryDate = loadEntries().values
            .compactMap { $0.completionDate }
            .max()

        if let latestEntryDate {
            return latestEntryDate
        }

        let legacy = UserDefaults.standard.double(forKey: legacyLastCheckInKey)
        guard legacy > 0 else { return nil }
        return Date(timeIntervalSince1970: legacy)
    }

    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

        // Current streak counts consecutive on-time months up to the reference month.
        var currentStreak = 0
        var cursor = referenceDate.startOfMonth
        while onTimeMonths.contains(cursor) {
            currentStreak += 1
            cursor = cursor.adding(months: -1).startOfMonth
        }

        // Best streak across history.
        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() {
        let defaults = UserDefaults.standard
        defaults.removeObject(forKey: notesKey)
        defaults.removeObject(forKey: completionsKey)
        defaults.removeObject(forKey: entriesKey)
        defaults.removeObject(forKey: legacyLastCheckInKey)
    }

    // MARK: - Private Helpers

    private static func updateEntry(for date: Date, mutate: (inout MonthlyCheckInEntry) -> Void) {
        let key = monthKey(for: date)
        var entries = loadEntries()
        var entry = entries[key] ?? MonthlyCheckInEntry(
            note: nil,
            rating: nil,
            mood: nil,
            completionTime: legacyCompletion(for: key),
            createdAt: Date().timeIntervalSince1970
        )

        mutate(&entry)

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

        let isEmpty = entry.note == nil && entry.rating == nil && entry.mood == nil && entry.completionTime == nil
        if isEmpty {
            entries.removeValue(forKey: key)
        } else {
            entries[key] = entry
        }

        saveEntries(entries)
        persistLegacyMirrors(entries)
    }

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

        // Ensure legacy data is merged if it existed before this release.
        return mergeLegacy(into: decoded)
    }

    private static func saveEntries(_ entries: [String: MonthlyCheckInEntry]) {
        guard let data = try? JSONEncoder().encode(entries) else { return }
        UserDefaults.standard.set(data, forKey: entriesKey)
    }

    private static func migrateLegacyData() -> [String: MonthlyCheckInEntry] {
        let notes = loadNotes()
        let completions = loadCompletions()
        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
            )
        }

        saveEntries(entries)
        return entries
    }

    private static func mergeLegacy(into entries: [String: MonthlyCheckInEntry]) -> [String: MonthlyCheckInEntry] {
        var merged = entries
        let notes = loadNotes()
        let completions = loadCompletions()
        var shouldSave = false

        for (key, note) in notes where merged[key]?.note == nil {
            var entry = merged[key] ?? MonthlyCheckInEntry(
                note: nil,
                rating: nil,
                mood: nil,
                completionTime: completions[key],
                createdAt: Date().timeIntervalSince1970
            )
            entry.note = note
            merged[key] = entry
            shouldSave = true
        }

        for (key, completion) in completions where merged[key]?.completionTime == nil {
            var entry = merged[key] ?? MonthlyCheckInEntry(
                note: nil,
                rating: nil,
                mood: nil,
                completionTime: nil,
                createdAt: completion
            )
            entry.completionTime = completion
            merged[key] = entry
            shouldSave = true
        }

        if shouldSave {
            saveEntries(merged)
        }

        return merged
    }

    private static func persistLegacyMirrors(_ entries: [String: MonthlyCheckInEntry]) {
        var notes: [String: String] = [:]
        var completions: [String: Double] = [:]

        for (key, entry) in entries {
            if let note = entry.note {
                notes[key] = note
            }
            if let completion = entry.completionTime {
                completions[key] = completion
            }
        }

        saveNotes(notes)
        saveCompletions(completions)
    }

    private static func loadNotes() -> [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 saveNotes(_ notes: [String: String]) {
        guard let data = try? JSONEncoder().encode(notes) else { return }
        UserDefaults.standard.set(data, forKey: notesKey)
    }

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

    private static func saveCompletions(_ completions: [String: Double]) {
        guard let data = try? JSONEncoder().encode(completions) else { return }
        UserDefaults.standard.set(data, forKey: completionsKey)
    }

    private static func legacyCompletion(for key: String) -> Double? {
        loadCompletions()[key]
    }

    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
    }
}
