gramora/gramora/Managers/SubscriptionManager.swift

import AppKit
import Combine
import Foundation
import StoreKit

enum PurchaseError: Equatable {
    case productsLoadTimeout
    case subscriptionUnavailable
    case planUnavailable
    case purchaseNotVerified
    case purchasePending
    case noActiveSubscriptions
    case generic

    var message: String {
        switch self {
        case .productsLoadTimeout:
            "Couldn't load subscription options right now. Check your internet connection and try again."
        case .subscriptionUnavailable:
            "This subscription is currently unavailable. Please try again in a moment."
        case .planUnavailable:
            "This plan isn't available right now. Please try again later."
        case .purchaseNotVerified:
            "Apple couldn't verify this purchase."
        case .purchasePending:
            "This purchase is waiting for approval (for example, Ask to Buy)."
        case .noActiveSubscriptions:
            "No active purchases were found for this Apple ID."
        case .generic:
            "Something went wrong with your subscription. Please try again."
        }
    }
}

@MainActor
final class SubscriptionManager: ObservableObject {
    static let maxFreeAIUses = 4
    private static let freeAIUsesConsumedKey = "freeAIUsesConsumed"

    private enum ProductLoadError: Error {
        case timeout
    }

    @Published private(set) var productsByID: [String: Product] = [:]
    @Published private(set) var isLoadingProducts = false
    @Published private(set) var purchasingPlan: PaywallPlan?
    @Published private(set) var hasPremiumAccess = false
    @Published private(set) var hasManageableSubscription = false
    @Published private(set) var hasResolvedPremiumStatus = false
    @Published private(set) var activePlan: PaywallPlan?
    @Published private(set) var activeRecurringPlan: PaywallPlan?
    @Published private(set) var introOfferEligibleByPlan: [PaywallPlan: Bool] = [:]
    @Published private(set) var freeAIUsesConsumed = 0
    @Published var purchaseError: PurchaseError?
    @Published var subscriptionSuccessMessage: String?

    var hasAllProductsLoaded: Bool {
        SubscriptionProductID.all.allSatisfy { productsByID[$0] != nil }
    }

    var remainingFreeAIUses: Int {
        max(0, Self.maxFreeAIUses - freeAIUsesConsumed)
    }

    /// Premium access confirmed by StoreKit after launch or purchase.
    var isPremiumActive: Bool {
        hasResolvedPremiumStatus && hasPremiumAccess
    }

    var canUseAIFeatures: Bool {
        isPremiumActive || remainingFreeAIUses > 0
    }

    private var inFlightProductLoad: Task<Void, Never>?
    private var cancellables = Set<AnyCancellable>()
    private let defaults: UserDefaults
    /// Keeps premium UI in sync after a verified purchase until StoreKit entitlements catch up.
    private var pendingVerifiedPurchase: (plan: PaywallPlan, purchasedAt: Date)?
    private static let pendingPurchaseGraceInterval: TimeInterval = 24 * 60 * 60

    init(defaults: UserDefaults = .standard) {
        self.defaults = defaults
        freeAIUsesConsumed = min(
            defaults.integer(forKey: Self.freeAIUsesConsumedKey),
            Self.maxFreeAIUses
        )

        observeAIFeatureUsage()

        Task { await listenForTransactionUpdates() }
        Task { await refreshPremiumAccess() }
        Task { await loadProducts(presentLoadingUI: false) }
        observeAppLifecycle()
    }

    func recordSuccessfulAIUse() {
        guard !isPremiumActive else { return }
        guard freeAIUsesConsumed < Self.maxFreeAIUses else { return }

        freeAIUsesConsumed += 1
        defaults.set(freeAIUsesConsumed, forKey: Self.freeAIUsesConsumedKey)
    }

    private func observeAIFeatureUsage() {
        NotificationCenter.default.publisher(for: .aiFeatureDidSucceed)
            .receive(on: DispatchQueue.main)
            .sink { [weak self] _ in
                self?.recordSuccessfulAIUse()
            }
            .store(in: &cancellables)
    }

    func loadProducts(presentLoadingUI: Bool? = nil) async {
        if let existing = inFlightProductLoad {
            await existing.value
            return
        }

        let uiFlag = presentLoadingUI
        let task = Task { @MainActor in
            await performLoadProducts(presentLoadingUI: uiFlag)
        }
        inFlightProductLoad = task
        await task.value
        inFlightProductLoad = nil
    }

    func product(for plan: PaywallPlan) -> Product? {
        productsByID[plan.productID]
    }

    func displayTitle(for plan: PaywallPlan) -> String {
        product(for: plan)?.displayName ?? plan.rawValue.capitalized
    }

    func mainPrice(for plan: PaywallPlan) -> String {
        product(for: plan)?.displayPrice ?? "—"
    }

    func secondaryPrice(for plan: PaywallPlan) -> String {
        guard let product = product(for: plan) else { return "—" }

        let formatStyle = product.priceFormatStyle

        switch plan {
        case .weekly:
            let doubled = product.price * 2
            return "\(doubled.formatted(formatStyle)) / week"
        case .monthly:
            let weekly = product.price / Decimal(4.33)
            return "\(weekly.formatted(formatStyle)) / week"
        case .yearly:
            let weekly = product.price / Decimal(52)
            return "\(weekly.formatted(formatStyle)) / week"
        }
    }

    func tagText(for plan: PaywallPlan) -> String? {
        if hasFreeTrialOffer(for: plan) {
            return "Free Trial"
        }

        switch plan {
        case .weekly:
            return "Basic"
        case .yearly:
            return savingsTagText(for: plan)
        case .monthly:
            return nil
        }
    }

    func ctaTitle(for plan: PaywallPlan) -> String {
        if hasFreeTrialOffer(for: plan) {
            return "START FOR FREE"
        }

        switch plan {
        case .weekly: return "START WEEKLY PLAN"
        case .monthly: return "START MONTHLY PLAN"
        case .yearly: return "START YEARLY PLAN"
        }
    }

    func billingDescription(for plan: PaywallPlan) -> String {
        guard let product = product(for: plan) else {
            return "Loading pricing…"
        }

        if hasFreeTrialOffer(for: plan),
           let subscription = product.subscription,
           let intro = subscription.introductoryOffer {
            let trialPeriod = formattedPeriod(intro.period)
            let billingPeriod = formattedSubscriptionPeriod(subscription.subscriptionPeriod)
            return "\(trialPeriod) Free Trial, then \(product.displayPrice) per \(billingPeriod)"
        }

        switch plan {
        case .weekly:
            return "Billed at \(product.displayPrice) every week"
        case .monthly:
            return "Billed at \(product.displayPrice) per month"
        case .yearly:
            return "Billed at \(product.displayPrice) every year"
        }
    }

    func showsStrikethroughSecondaryPrice(for plan: PaywallPlan) -> Bool {
        plan == .weekly
    }

    private func savingsTagText(for plan: PaywallPlan) -> String? {
        guard plan == .yearly,
              let yearlyProduct = product(for: .yearly),
              let weeklyProduct = product(for: .weekly),
              weeklyProduct.price > 0 else { return nil }

        let yearlyWeeklyEquivalent = yearlyProduct.price / Decimal(52)
        let savings = (1 - (yearlyWeeklyEquivalent / weeklyProduct.price)) * 100
        let rounded = Int(NSDecimalNumber(decimal: savings).doubleValue.rounded())
        guard rounded > 0 else { return nil }
        return "Save \(rounded)%"
    }

    func hasFreeTrialOffer(for plan: PaywallPlan) -> Bool {
        guard introOfferEligibleByPlan[plan] == true else { return false }
        guard let product = product(for: plan),
              let subscription = product.subscription,
              let intro = subscription.introductoryOffer,
              intro.paymentMode == .freeTrial else { return false }
        return true
    }

    func purchase(_ plan: PaywallPlan) async -> Bool {
        purchaseError = nil

        if isPremiumActive {
            if !hasManageableSubscription {
                return true
            }
            if let activePlan, plan == activePlan {
                return true
            }
            if let activeRecurringPlan, plan == activeRecurringPlan {
                return true
            }
        }

        var product = productsByID[plan.productID]
        if product == nil {
            await loadProducts(presentLoadingUI: false)
            product = productsByID[plan.productID]
        }

        guard let product else {
            purchaseError = productsByID.isEmpty ? .productsLoadTimeout : .planUnavailable
            #if DEBUG
            let missingProductIDs = SubscriptionProductID.all.filter { productsByID[$0] == nil }
            if !missingProductIDs.isEmpty {
                print(
                    """
                    [SubscriptionManager] Missing StoreKit products: \(missingProductIDs.joined(separator: ", "))
                    To test locally, attach `Paywall.storekit` in Xcode:
                    Product > Scheme > Edit Scheme > Run > Options > StoreKit Configuration.
                    """
                )
            }
            #endif
            return false
        }

        purchasingPlan = plan
        defer { purchasingPlan = nil }

        do {
            let result = try await product.purchase()
            switch result {
            case .success(let verification):
                guard case .verified(let transaction) = verification else {
                    purchaseError = .purchaseNotVerified
                    return false
                }
                await transaction.finish()
                applyOptimisticPremiumAccess(for: transaction)
                await refreshPremiumAccessWithRetry()
                if hasPremiumAccess {
                    NotificationCenter.default.post(name: .subscriptionPurchased, object: nil)
                } else {
                    subscriptionSuccessMessage = "Purchase successful. If premium features don't unlock right away, tap Restore Purchases."
                }
                return hasPremiumAccess
            case .userCancelled:
                return false
            case .pending:
                purchaseError = .purchasePending
                return false
            @unknown default:
                purchaseError = .generic
                return false
            }
        } catch {
            purchaseError = .generic
            return false
        }
    }

    func openSubscriptionManagement() {
        guard let url = URL(string: "https://apps.apple.com/account/subscriptions") else { return }
        NSWorkspace.shared.open(url)
    }

    func restorePurchases() async {
        purchaseError = nil
        subscriptionSuccessMessage = nil
        let hadPremium = isPremiumActive
        do {
            try await AppStore.sync()
            await refreshPremiumAccess()
            if hasPremiumAccess {
                if !hadPremium {
                    subscriptionSuccessMessage = "Your purchases have been restored."
                }
            } else {
                purchaseError = .noActiveSubscriptions
            }
        } catch {
            purchaseError = .generic
        }
    }

    func clearSubscriptionSuccessMessage() {
        subscriptionSuccessMessage = nil
    }

    private func performLoadProducts(presentLoadingUI: Bool?) async {
        let missingAny = SubscriptionProductID.all.contains { productsByID[$0] == nil }
        let shouldShowLoading = presentLoadingUI ?? missingAny

        if !missingAny, presentLoadingUI != true {
            purchaseError = nil
            await refreshIntroOfferEligibility()
            await refreshPremiumAccess()
            return
        }

        if shouldShowLoading {
            isLoadingProducts = true
        }
        defer {
            if shouldShowLoading {
                isLoadingProducts = false
            }
        }

        purchaseError = nil
        do {
            let loaded = try await loadProductsWithTimeout(seconds: 12)
            var map: [String: Product] = [:]
            for product in loaded {
                map[product.id] = product
            }
            productsByID = map
            await refreshIntroOfferEligibility()
            await refreshPremiumAccess()
        } catch is ProductLoadError {
            if presentLoadingUI == true {
                purchaseError = .productsLoadTimeout
            }
        } catch {
            if presentLoadingUI == true {
                purchaseError = .generic
            }
        }
    }

    private func loadProductsWithTimeout(seconds: TimeInterval) async throws -> [Product] {
        try await withThrowingTaskGroup(of: [Product].self) { group in
            group.addTask {
                try await Product.products(for: SubscriptionProductID.all)
            }
            group.addTask {
                let timeoutNanoseconds = UInt64(seconds * 1_000_000_000)
                try await Task.sleep(nanoseconds: timeoutNanoseconds)
                throw ProductLoadError.timeout
            }

            guard let first = try await group.next() else {
                throw ProductLoadError.timeout
            }
            group.cancelAll()
            return first
        }
    }

    private func refreshPremiumAccess() async {
        var hasPremium = false
        var manageableSubscription = false
        var resolvedPlan: PaywallPlan?
        var resolvedRecurringPlan: PaywallPlan?
        var revokedPendingPlan = false

        for await result in Transaction.currentEntitlements {
            guard case .verified(let transaction) = result else { continue }
            guard SubscriptionProductID.all.contains(transaction.productID) else { continue }

            if transaction.revocationDate != nil {
                if pendingVerifiedPurchase?.plan == SubscriptionProductID.plan(for: transaction.productID) {
                    revokedPendingPlan = true
                    pendingVerifiedPurchase = nil
                }
                continue
            }

            guard let plan = SubscriptionProductID.plan(for: transaction.productID) else { continue }
            guard await isRecurringSubscriptionEntitled(transaction: transaction) else {
                continue
            }

            manageableSubscription = true
            hasPremium = true
            pendingVerifiedPurchase = nil
            if resolvedRecurringPlan == nil {
                resolvedRecurringPlan = plan
            }
            if resolvedPlan == nil {
                resolvedPlan = plan
            }
        }

        if !hasPremium,
           !revokedPendingPlan,
           let pending = pendingVerifiedPurchase,
           Date().timeIntervalSince(pending.purchasedAt) < Self.pendingPurchaseGraceInterval {
            hasPremium = true
            manageableSubscription = true
            resolvedPlan = pending.plan
            resolvedRecurringPlan = pending.plan
        }

        hasPremiumAccess = hasPremium
        hasManageableSubscription = manageableSubscription
        activePlan = resolvedPlan
        activeRecurringPlan = resolvedRecurringPlan
        hasResolvedPremiumStatus = true
    }

    private func applyOptimisticPremiumAccess(for transaction: Transaction) {
        guard SubscriptionProductID.all.contains(transaction.productID),
              transaction.revocationDate == nil,
              let plan = SubscriptionProductID.plan(for: transaction.productID) else { return }

        pendingVerifiedPurchase = (plan: plan, purchasedAt: Date())
        hasPremiumAccess = true
        hasResolvedPremiumStatus = true
        activePlan = plan
        hasManageableSubscription = true
        activeRecurringPlan = plan
    }

    private func refreshIntroOfferEligibility() async {
        var eligibility: [PaywallPlan: Bool] = [:]

        for plan in PaywallPlan.allCases {
            guard let product = productsByID[plan.productID],
                  let subscription = product.subscription,
                  let intro = subscription.introductoryOffer,
                  intro.paymentMode == .freeTrial else { continue }

            eligibility[plan] = await subscription.isEligibleForIntroOffer
        }

        introOfferEligibleByPlan = eligibility
    }

    private func refreshPremiumAccessWithRetry() async {
        await refreshPremiumAccess()
        guard !hasPremiumAccess else { return }

        for _ in 0..<6 {
            try? await Task.sleep(nanoseconds: 500_000_000)
            await refreshPremiumAccess()
            if hasPremiumAccess { return }
        }

        if let pending = pendingVerifiedPurchase,
           Date().timeIntervalSince(pending.purchasedAt) < Self.pendingPurchaseGraceInterval {
            hasPremiumAccess = true
            hasResolvedPremiumStatus = true
            hasManageableSubscription = true
            activePlan = pending.plan
            activeRecurringPlan = pending.plan
        }
    }

    private func isRecurringSubscriptionEntitled(transaction: Transaction) async -> Bool {
        guard transaction.revocationDate == nil else { return false }

        let productID = transaction.productID
        let product: Product?
        if let cached = productsByID[productID] {
            product = cached
        } else if let loaded = try? await Product.products(for: [productID]).first {
            product = loaded
        } else {
            product = nil
        }

        if let product, let subscription = product.subscription {
            do {
                let statuses = try await subscription.status
                if statuses.contains(where: {
                    switch $0.state {
                    case .subscribed, .inGracePeriod, .inBillingRetryPeriod:
                        return true
                    case .expired, .revoked:
                        return false
                    default:
                        return false
                    }
                }) {
                    return true
                }
                if statuses.contains(where: { $0.state == .expired || $0.state == .revoked }) {
                    return false
                }
            } catch {
                // Fall through to expiration-date check when status is unavailable.
            }
        }

        // Verified entitlement with a future expiration is still active when metadata/status is unreachable.
        if let expirationDate = transaction.expirationDate {
            return expirationDate > Date()
        }

        return false
    }

    private func observeAppLifecycle() {
        NotificationCenter.default.publisher(for: NSApplication.didBecomeActiveNotification)
            .sink { [weak self] _ in
                Task { await self?.refreshPremiumAccess() }
            }
            .store(in: &cancellables)
    }

    private func listenForTransactionUpdates() async {
        for await update in Transaction.updates {
            switch update {
            case .verified(let transaction):
                if transaction.revocationDate != nil,
                   let plan = SubscriptionProductID.plan(for: transaction.productID),
                   pendingVerifiedPurchase?.plan == plan {
                    pendingVerifiedPurchase = nil
                }
                await transaction.finish()
                await refreshPremiumAccess()
            case .unverified(_, let error):
                #if DEBUG
                print("[SubscriptionManager] Unverified transaction update: \(error)")
                #endif
                purchaseError = .purchaseNotVerified
            }
        }
    }

    private func formattedPeriod(_ period: Product.SubscriptionPeriod) -> String {
        switch period.unit {
        case .day where period.value == 1: "1 Day"
        case .day: "\(period.value) Days"
        case .week where period.value == 1: "1 Week"
        case .week: "\(period.value) Weeks"
        case .month where period.value == 1: "1 Month"
        case .month: "\(period.value) Months"
        case .year where period.value == 1: "1 Year"
        case .year: "\(period.value) Years"
        @unknown default: "\(period.value) Days"
        }
    }

    private func formattedSubscriptionPeriod(_ period: Product.SubscriptionPeriod) -> String {
        switch period.unit {
        case .day: period.value == 1 ? "day" : "\(period.value) days"
        case .week: period.value == 1 ? "week" : "\(period.value) weeks"
        case .month: period.value == 1 ? "month" : "\(period.value) months"
        case .year: period.value == 1 ? "year" : "\(period.value) years"
        @unknown default: "billing period"
        }
    }
}