app_for_reddit/App AI for Reddit/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 noActivePurchases
    case restoreSyncTimeout
    case generic
}

@MainActor
final class SubscriptionManager: ObservableObject {
    private enum ProductLoadError: Error {
        case timeout
    }

    private enum StoreSyncError: Error {
        case timeout
    }

    private static let cachedPremiumKey = "SubscriptionManager.cachedIsPremium"
    private static let cachedPremiumPlanKey = "SubscriptionManager.cachedPremiumPlan"
    private static let cachedHasEverPurchasedKey = "SubscriptionManager.cachedHasEverPurchased"

    @Published private(set) var productsByID: [String: Product] = [:]
    @Published private(set) var isLoadingProducts = false
    @Published private(set) var purchasingPlan: PaywallPlan?
    @Published private(set) var isRestoringPurchases = false
    @Published private(set) var hasPremiumAccess = false
    @Published private(set) var hasResolvedPremiumStatus = false
    @Published private(set) var activePremiumPlan: PaywallPlan?
    @Published private(set) var productLoadError: String?
    @Published private(set) var introOfferEligibleByProductID: [String: Bool] = [:]
    @Published private(set) var hasEverPurchasedPremium = false
    @Published var purchaseError: PurchaseError?

    var availablePaywallPlans: [PaywallPlan] {
        if hasEverPurchasedPremium {
            return PaywallPlan.allCases.filter { $0 != .lifetime }
        }
        return PaywallPlan.allCases
    }

    var isResolvingEntitlements: Bool {
        !hasResolvedPremiumStatus
    }

    var hasLifetimeAccess: Bool {
        activePremiumPlan == .lifetime
    }

    var hasActiveSubscription: Bool {
        activePremiumPlan == .monthly || activePremiumPlan == .yearly
    }

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

    private var inFlightProductLoad: Task<Void, Never>?
    private var initialEntitlementTask: Task<Void, Never>?
    private var hasStarted = false

    init() {
        hasPremiumAccess = UserDefaults.standard.bool(forKey: Self.cachedPremiumKey)
        hasEverPurchasedPremium = UserDefaults.standard.bool(forKey: Self.cachedHasEverPurchasedKey)
        if let raw = UserDefaults.standard.string(forKey: Self.cachedPremiumPlanKey),
           let plan = PaywallPlan(rawValue: raw) {
            activePremiumPlan = plan
        }
    }

    func start() {
        guard !hasStarted else { return }
        hasStarted = true

        Task { await listenForTransactionUpdates() }

        initialEntitlementTask = Task {
            await refreshPremiumAccess()
            await loadProducts(presentLoadingUI: false)
        }
    }

    func ensureEntitlementsResolved() async {
        await initialEntitlementTask?.value
    }

    func eligibleIntroOffer(for plan: PaywallPlan) -> Product.SubscriptionOffer? {
        guard let product = product(for: plan) else { return nil }
        return eligibleIntroOffer(for: product)
    }

    func eligibleIntroOffer(for product: Product) -> Product.SubscriptionOffer? {
        guard !hasEverPurchasedPremium,
              introOfferEligibleByProductID[product.id] == true,
              let offer = product.subscription?.introductoryOffer,
              offer.paymentMode == .freeTrial else {
            return nil
        }
        return offer
    }

    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 mainPrice(for plan: PaywallPlan) -> String {
        let config = PaywallConfigService.shared.config
        return product(for: plan)?.displayPrice ?? plan.planCopy(from: config).fallbackPrice
    }

    func priceSuffix(for plan: PaywallPlan) -> String? {
        plan.planCopy(from: PaywallConfigService.shared.config).priceSuffix
    }

    func billingDescription(for plan: PaywallPlan) -> String {
        let config = PaywallConfigService.shared.config
        guard let product = product(for: plan) else {
            return plan.fallbackBillingDescription(config: config)
        }

        if !hasEverPurchasedPremium,
           let subscription = product.subscription,
           let intro = subscription.introductoryOffer,
           intro.paymentMode == .freeTrial,
           introOfferEligibleByProductID[product.id] == true {
            let trialPeriod = formattedPeriod(intro.period)
            let billingPeriod = formattedSubscriptionPeriod(subscription.subscriptionPeriod)
            return "\(trialPeriod) Free Trial, then \(product.displayPrice) per \(billingPeriod)"
        }

        switch plan {
        case .monthly:
            return "Billed at \(product.displayPrice) every month"
        case .yearly:
            return "Billed at \(product.displayPrice) every year"
        case .lifetime:
            return "One-time payment of \(product.displayPrice)"
        }
    }

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

        guard availablePaywallPlans.contains(plan) else {
            purchaseError = .planUnavailable
            return false
        }

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

        guard let product else {
            let missingProductIDs = SubscriptionProductID.all.filter { productsByID[$0] == nil }
            if missingProductIDs.isEmpty {
                purchaseError = .subscriptionUnavailable
            } else {
                purchaseError = .planUnavailable
                #if DEBUG
                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
                }

                let trustedPlan = planGrantingPremium(from: transaction)
                if let trustedPlan {
                    applyPremiumAccess(from: transaction)
                }

                await transaction.finish()
                await reconcilePremiumAccess(trustedPlan: trustedPlan)
                return hasPremiumAccess
            case .userCancelled:
                return false
            case .pending:
                purchaseError = .purchasePending
                return false
            @unknown default:
                return false
            }
        } catch {
            purchaseError = .generic
            return false
        }
    }

    func openSubscriptionManagement() {
        let urlString = PaywallConfigService.shared.config.urls.manageSubscriptions
        guard let url = URL(string: urlString) else { return }
        NSWorkspace.shared.open(url)
    }

    func restorePurchases() async {
        purchaseError = nil
        isRestoringPurchases = true
        defer { isRestoringPurchases = false }

        do {
            try await syncAppStoreWithTimeout(seconds: 30)
        } catch is StoreSyncError {
            purchaseError = .restoreSyncTimeout
            return
        } catch {
            purchaseError = .generic
            return
        }

        await reconcilePremiumAccess(trustedPlan: nil)
        if !hasPremiumAccess {
            purchaseError = .noActivePurchases
        }
    }

    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 refreshPremiumAccess()
            return
        }

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

        purchaseError = nil
        productLoadError = nil
        do {
            let loaded = try await loadProductsWithTimeout(seconds: 12)
            guard !loaded.isEmpty else {
                productLoadError = PaywallConfigService.shared.config.messages.noPlansAvailable
                productsByID = [:]
                return
            }
            var map: [String: Product] = [:]
            for product in loaded {
                map[product.id] = product
            }
            productsByID = map
            productLoadError = nil
            await refreshIntroOfferEligibility()
            await refreshPremiumAccess()
        } catch is ProductLoadError {
            productLoadError = PaywallConfigService.shared.config.messages.plansLoadFailed
            purchaseError = .productsLoadTimeout
        } catch {
            productLoadError = PaywallConfigService.shared.config.messages.plansLoadFailed
            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(trustedPlan: PaywallPlan? = nil) async {
        var resolvedPlan: PaywallPlan?

        if let lifetimePlan = await activeLifetimePlan() {
            resolvedPlan = lifetimePlan
        }

        if resolvedPlan == nil {
            for await result in Transaction.currentEntitlements {
                guard case .verified(let transaction) = result else { continue }
                guard let plan = planGrantingPremium(from: transaction) else { continue }
                resolvedPlan = preferredPlan(resolvedPlan, plan)
            }
        }

        if resolvedPlan == nil || resolvedPlan != .lifetime {
            if let subscriptionPlan = await activeSubscriptionPlan() {
                resolvedPlan = preferredPlan(resolvedPlan, subscriptionPlan)
            }
        }

        let finalPlan = mergeTrustedPlan(resolvedPlan, trustedPlan)
        activePremiumPlan = finalPlan
        hasPremiumAccess = finalPlan != nil
        hasResolvedPremiumStatus = true
        persistPremiumCache()
        await refreshPurchaseHistory()
    }

    /// Re-reads entitlements from StoreKit, then syncs with the App Store when needed.
    /// A verified purchase is never downgraded while the network is still catching up.
    private func reconcilePremiumAccess(trustedPlan: PaywallPlan?) async {
        await refreshPremiumAccess(trustedPlan: trustedPlan)

        guard let trustedPlan, !hasPremiumAccess else { return }

        for attempt in 1...3 {
            do {
                try await syncAppStoreWithTimeout(seconds: 15)
            } catch {
                continue
            }

            await refreshPremiumAccess(trustedPlan: trustedPlan)
            if hasPremiumAccess { return }

            let delayNanoseconds = UInt64(attempt) * 1_500_000_000
            try? await Task.sleep(nanoseconds: delayNanoseconds)
        }

        applyTrustedPlan(trustedPlan)
    }

    private func mergeTrustedPlan(_ resolvedPlan: PaywallPlan?, _ trustedPlan: PaywallPlan?) -> PaywallPlan? {
        switch (resolvedPlan, trustedPlan) {
        case (nil, nil):
            nil
        case (let resolved?, nil):
            resolved
        case (nil, let trusted?):
            trusted
        case (let resolved?, let trusted?):
            preferredPlan(resolved, trusted)
        }
    }

    private func applyTrustedPlan(_ plan: PaywallPlan) {
        activePremiumPlan = preferredPlan(activePremiumPlan, plan)
        hasPremiumAccess = true
        hasResolvedPremiumStatus = true
        markAsHavingPurchasedPremium()
        persistPremiumCache()
    }

    private func syncAppStoreWithTimeout(seconds: TimeInterval) async throws {
        try await withThrowingTaskGroup(of: Void.self) { group in
            group.addTask {
                try await AppStore.sync()
            }
            group.addTask {
                let timeoutNanoseconds = UInt64(seconds * 1_000_000_000)
                try await Task.sleep(nanoseconds: timeoutNanoseconds)
                throw StoreSyncError.timeout
            }

            try await group.next()
            group.cancelAll()
        }
    }

    private func applyPremiumAccess(from transaction: StoreKit.Transaction) {
        guard let plan = planGrantingPremium(from: transaction) else { return }
        activePremiumPlan = preferredPlan(activePremiumPlan, plan)
        hasPremiumAccess = true
        hasResolvedPremiumStatus = true
        markAsHavingPurchasedPremium()
        persistPremiumCache()
    }

    private func markAsHavingPurchasedPremium() {
        guard !hasEverPurchasedPremium else { return }
        hasEverPurchasedPremium = true
        UserDefaults.standard.set(true, forKey: Self.cachedHasEverPurchasedKey)
    }

    private func refreshPurchaseHistory() async {
        if hasEverPurchasedPremium { return }

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

    private func persistPremiumCache() {
        UserDefaults.standard.set(hasPremiumAccess, forKey: Self.cachedPremiumKey)
        UserDefaults.standard.set(activePremiumPlan?.rawValue, forKey: Self.cachedPremiumPlanKey)
    }

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

        for product in productsByID.values {
            guard let subscription = product.subscription,
                  subscription.introductoryOffer != nil else {
                continue
            }
            let groupID = subscription.subscriptionGroupID
            eligibility[product.id] = await Product.SubscriptionInfo.isEligibleForIntroOffer(for: groupID)
        }

        introOfferEligibleByProductID = eligibility
    }

    private func planGrantingPremium(from transaction: StoreKit.Transaction) -> PaywallPlan? {
        guard transactionGrantsPremium(transaction) else { return nil }
        return SubscriptionProductID.plan(for: transaction.productID)
    }

    private func transactionGrantsPremium(_ transaction: StoreKit.Transaction) -> Bool {
        guard SubscriptionProductID.all.contains(transaction.productID) else { return false }
        guard transaction.revocationDate == nil else { return false }
        if let expirationDate = transaction.expirationDate, expirationDate < Date() {
            return false
        }
        return true
    }

    private func preferredPlan(_ current: PaywallPlan?, _ candidate: PaywallPlan) -> PaywallPlan {
        if current == .lifetime || candidate == .lifetime {
            return .lifetime
        }
        if let current {
            if current == .yearly || candidate == .yearly {
                return .yearly
            }
            return .monthly
        }
        return candidate
    }

    private func activeLifetimePlan() async -> PaywallPlan? {
        guard let result = await Transaction.latest(for: SubscriptionProductID.lifetime) else {
            return nil
        }
        guard case .verified(let transaction) = result else { return nil }
        return planGrantingPremium(from: transaction)
    }

    private func activeSubscriptionPlan() async -> PaywallPlan? {
        var resolvedPlan: PaywallPlan?

        for await result in Transaction.currentEntitlements {
            guard case .verified(let transaction) = result else { continue }
            guard let plan = planGrantingPremium(from: transaction), plan != .lifetime else { continue }
            resolvedPlan = preferredPlan(resolvedPlan, plan)
        }

        if resolvedPlan == nil, await hasActiveSubscriptionStatus() {
            resolvedPlan = await subscriptionPlanFromProductStatus()
        }

        return resolvedPlan
    }

    private func subscriptionPlanFromProductStatus() async -> PaywallPlan? {
        if productsByID.isEmpty {
            do {
                let loaded = try await Product.products(for: SubscriptionProductID.all)
                var map: [String: Product] = [:]
                for product in loaded { map[product.id] = product }
                productsByID = map
            } catch {
                return nil
            }
        }

        for product in productsByID.values {
            guard let subscription = product.subscription else { continue }
            guard let plan = SubscriptionProductID.plan(for: product.id), plan != .lifetime else { continue }

            do {
                let statuses = try await subscription.status
                let isActive = statuses.contains { status in
                    switch status.state {
                    case .subscribed, .inGracePeriod, .inBillingRetryPeriod:
                        return true
                    case .expired, .revoked:
                        return false
                    default:
                        return false
                    }
                }
                if isActive { return plan }
            } catch {
                continue
            }
        }

        return nil
    }

    private func hasActiveSubscriptionStatus() async -> Bool {
        await subscriptionPlanFromProductStatus() != nil
    }

    private func listenForTransactionUpdates() async {
        for await update in Transaction.updates {
            guard case .verified(let transaction) = update else { continue }

            let trustedPlan = planGrantingPremium(from: transaction)
            if let trustedPlan {
                applyPremiumAccess(from: transaction)
            }

            await transaction.finish()
            await reconcilePremiumAccess(trustedPlan: trustedPlan)
        }
    }

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