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 availablePaywallPlans: [PaywallPlan] = []
    @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 hasEverPurchasedPremium = false
    @Published private(set) var trialDisplayByPlan: [PaywallPlan: PaywallTrialDisplay] = [:]
    @Published var purchaseError: PurchaseError?

    var isResolvingEntitlements: Bool {
        !hasResolvedPremiumStatus
    }

    var hasLifetimeAccess: Bool {
        activePremiumPlan?.isLifetime == true
    }

    var hasActiveSubscription: Bool {
        guard let activePremiumPlan else { return false }
        return !activePremiumPlan.isLifetime
    }

    var hasAllProductsLoaded: Bool {
        !availablePaywallPlans.isEmpty
    }

    var isEligibleForFreeAITier: Bool {
        !hasEverPurchasedPremium
    }

    var isEligibleForIntroTrial: Bool {
        !hasEverPurchasedPremium && !hasPremiumAccess
    }

    private var inFlightProductLoad: Task<Void, Never>?
    private var initialEntitlementTask: Task<Void, Never>?
    private var entitlementRefreshTask: Task<Void, Never>?
    private var hasStarted = false
    private var launchedWithCachedPremium = false
    private var cancellables = Set<AnyCancellable>()
    private var lastBackgroundEntitlementRefresh: Date?

    /// Minimum time between entitlement re-checks when the app regains focus.
    private static let backgroundRefreshInterval: TimeInterval = 120

    init() {
        hasEverPurchasedPremium = UserDefaults.standard.bool(forKey: Self.cachedHasEverPurchasedKey)

        // Optimistic restore so returning subscribers aren't paywalled while StoreKit resolves.
        if UserDefaults.standard.bool(forKey: Self.cachedPremiumKey) {
            launchedWithCachedPremium = true
            hasPremiumAccess = true
            if let planRaw = UserDefaults.standard.string(forKey: Self.cachedPremiumPlanKey),
               let plan = PaywallPlan.from(rawValue: planRaw) {
                activePremiumPlan = plan
            }
        }
    }

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

        observeAppLifecycle()
        Task { await listenForTransactionUpdates() }

        initialEntitlementTask = Task {
            // Local StoreKit entitlements only on launch. AppStore.sync() prompts for
            // Apple ID and fires on every relaunch when purchase flags are cached.
            await refreshPremiumAccess()

            if launchedWithCachedPremium {
                launchedWithCachedPremium = false
                hasResolvedPremiumStatus = true
                persistPremiumCache()
            }

            await loadProducts(presentLoadingUI: false)
        }
    }

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

    private func refreshEntitlementsOnAppActive() {
        guard hasStarted else { return }

        if let lastRefresh = lastBackgroundEntitlementRefresh,
           Date().timeIntervalSince(lastRefresh) < Self.backgroundRefreshInterval {
            return
        }

        entitlementRefreshTask?.cancel()
        entitlementRefreshTask = Task {
            // Local StoreKit entitlements only — AppStore.sync() prompts for Apple ID
            // and can loop in Xcode when focus returns after dismissing the dialog.
            await refreshPremiumAccess()
            lastBackgroundEntitlementRefresh = Date()
        }
    }

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

    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 trialDisplay(for plan: PaywallPlan) -> PaywallTrialDisplay? {
        trialDisplayByPlan[plan]
    }

    func refreshTrialEligibility() async {
        await refreshTrialOffers()
    }

    func mainPrice(for plan: PaywallPlan) -> String {
        let config = PaywallConfigService.shared.config
        return product(for: plan)?.displayPrice ?? config.loadingPrice
    }

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

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

        if product.type == .nonConsumable {
            return "One-time payment of \(product.displayPrice)"
        }

        let period = product.subscription.map {
            config.subscriptionPeriodDescription(from: $0.subscriptionPeriod)
        } ?? config.messages.periodFallback
        return "Billed at \(product.displayPrice) every \(period)"
    }

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

        guard availablePaywallPlans.contains(where: { $0 == 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 {
            purchaseError = .planUnavailable
            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 }
        RedditWebViewStore.shared.suspendForSubscriptionManagement()
        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 config = PaywallConfigService.shared.config
        let requestedProductIDs = Set(SubscriptionProductID.productIDs(from: config))
        let missingAny = availablePaywallPlans.isEmpty
        let shouldShowLoading = presentLoadingUI ?? missingAny

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

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

        purchaseError = nil
        productLoadError = nil
        do {
            let loaded = try await loadProductsWithTimeout(
                productIDs: requestedProductIDs,
                seconds: 12
            )
            guard !loaded.isEmpty else {
                productLoadError = PaywallConfigService.shared.config.messages.noPlansAvailable
                productsByID = [:]
                availablePaywallPlans = []
                return
            }
            var map: [String: Product] = [:]
            for product in loaded {
                map[product.id] = product
            }
            productsByID = map
            updateAvailablePaywallPlans(from: map, requestedProductIDs: requestedProductIDs)
            productLoadError = nil
            await refreshPremiumAccess()
            await refreshTrialOffers()
        } catch is ProductLoadError {
            productLoadError = PaywallConfigService.shared.config.messages.plansLoadFailed
            purchaseError = .productsLoadTimeout
        } catch {
            productLoadError = PaywallConfigService.shared.config.messages.plansLoadFailed
            purchaseError = .generic
        }
    }

    private func updateAvailablePaywallPlans(
        from productsByID: [String: Product],
        requestedProductIDs: Set<String>
    ) {
        availablePaywallPlans = requestedProductIDs
            .compactMap { productID in
                productsByID[productID].map { _ in PaywallPlan(productID: productID) }
            }
            .sorted(by: PaywallPlan.sortOrder)
    }

    private func loadProductsWithTimeout(productIDs: Set<String>, seconds: TimeInterval) async throws -> [Product] {
        try await withThrowingTaskGroup(of: [Product].self) { group in
            group.addTask {
                try await Product.products(for: productIDs)
            }
            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?.isLifetime != true {
            if let subscriptionPlan = await activeSubscriptionPlan() {
                resolvedPlan = preferredPlan(resolvedPlan, subscriptionPlan)
            }
        }

        if resolvedPlan == nil {
            resolvedPlan = await activePlanFromLatestTransactions()
        }

        let finalPlan = mergeTrustedPlan(resolvedPlan, trustedPlan)

        if let finalPlan {
            activePremiumPlan = finalPlan
            hasPremiumAccess = true
            launchedWithCachedPremium = false
            markAsHavingPurchasedPremium()
        } else if launchedWithCachedPremium, activePremiumPlan != nil {
            // StoreKit can lag behind on launch; keep the cached plan until reconciliation finishes.
            hasPremiumAccess = true
        } else {
            activePremiumPlan = nil
            hasPremiumAccess = false
        }

        let resolved = !launchedWithCachedPremium
        if resolved {
            hasResolvedPremiumStatus = true
            persistPremiumCache()
            await syncPurchaseHistory()
        }
        await refreshTrialOffers()
    }

    /// Re-reads entitlements from StoreKit, then syncs with the App Store when a
    /// purchase was just verified but entitlements have not caught up yet.
    private func reconcilePremiumAccess(trustedPlan: PaywallPlan?) async {
        await refreshPremiumAccess(trustedPlan: trustedPlan)

        guard let trustedPlan, !hasPremiumAccess else { return }

        do {
            try await syncAppStoreWithTimeout(seconds: 15)
        } catch {
            applyTrustedPlan(trustedPlan)
            return
        }

        await refreshPremiumAccess(trustedPlan: trustedPlan)
        if !hasPremiumAccess {
            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) {
        launchedWithCachedPremium = false
        activePremiumPlan = preferredPlan(activePremiumPlan, plan)
        hasPremiumAccess = true
        hasResolvedPremiumStatus = true
        markAsHavingPurchasedPremium()
        persistPremiumCache()
    }

    private func revokePremiumAccess() {
        launchedWithCachedPremium = false
        activePremiumPlan = nil
        hasPremiumAccess = false
        hasResolvedPremiumStatus = true
        markAsHavingPurchasedPremium()
        AIFreeUsageManager.shared.forfeitFreeUses()
        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 }
        launchedWithCachedPremium = false
        activePremiumPlan = preferredPlan(activePremiumPlan, plan)
        hasPremiumAccess = true
        hasResolvedPremiumStatus = true
        markAsHavingPurchasedPremium()
        persistPremiumCache()
    }

    private func markAsHavingPurchasedPremium() {
        let wasAlreadyMarked = hasEverPurchasedPremium
        hasEverPurchasedPremium = true
        trialDisplayByPlan = [:]
        UserDefaults.standard.set(true, forKey: Self.cachedHasEverPurchasedKey)
        if !wasAlreadyMarked {
            AIFreeUsageManager.shared.forfeitFreeUses()
        }
    }

    /// Syncs `hasEverPurchasedPremium` with StoreKit history.
    /// - Refund: revoked transactions remain in history → flag stays `true` → no trial badge.
    /// - Sandbox reset: no transactions → flag clears → trial badge returns for testing.
    private func syncPurchaseHistory() async {
        var foundPurchase = false

        for await result in Transaction.all {
            guard case .verified(let transaction) = result else { continue }
            guard SubscriptionProductID.productIDs(from: PaywallConfigService.shared.config).contains(transaction.productID) else { continue }

            foundPurchase = true
            if !hasEverPurchasedPremium {
                hasEverPurchasedPremium = true
                UserDefaults.standard.set(true, forKey: Self.cachedHasEverPurchasedKey)
                AIFreeUsageManager.shared.forfeitFreeUses()
            }
            return
        }

        if !foundPurchase, hasEverPurchasedPremium {
            hasEverPurchasedPremium = false
            UserDefaults.standard.set(false, forKey: Self.cachedHasEverPurchasedKey)
        }
    }

    private func refreshTrialOffers() async {
        // Trial badge is first-time users only — never after a purchase or refund.
        guard isEligibleForIntroTrial else {
            trialDisplayByPlan = [:]
            return
        }

        guard let eligiblePlan = PaywallConfigService.shared.config.trialEligiblePlan else {
            trialDisplayByPlan = [:]
            return
        }

        if let display = await resolveTrialDisplay(for: eligiblePlan) {
            trialDisplayByPlan = [eligiblePlan: display]
        } else {
            trialDisplayByPlan = [:]
        }
    }

    private func resolveTrialDisplay(for plan: PaywallPlan) async -> PaywallTrialDisplay? {
        let config = PaywallConfigService.shared.config
        guard let eligiblePlan = config.trialEligiblePlan,
              plan.productID == eligiblePlan.productID else { return nil }
        guard let fallback = config.trial.fallbackDuration else { return nil }

        if let product = product(for: plan),
           let subscription = product.subscription,
           await subscription.isEligibleForIntroOffer,
           let offer = subscription.introductoryOffer {
            return PaywallTrialDisplay.from(offer: offer, config: config)
        }

        return PaywallTrialDisplay.from(fallback: fallback, config: config)
    }

    private func persistPremiumCache() {
        UserDefaults.standard.set(hasPremiumAccess, forKey: Self.cachedPremiumKey)
        if let plan = activePremiumPlan {
            UserDefaults.standard.set(plan.productID, forKey: Self.cachedPremiumPlanKey)
        } else {
            UserDefaults.standard.removeObject(forKey: Self.cachedPremiumPlanKey)
        }
    }

    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 {
        let knownProductIDs = Set(SubscriptionProductID.productIDs(from: PaywallConfigService.shared.config))
        guard knownProductIDs.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 {
        let order = [
            SubscriptionProductID.lifetime,
            SubscriptionProductID.yearly,
            SubscriptionProductID.monthly,
        ]

        func rank(for plan: PaywallPlan) -> Int {
            order.firstIndex(of: plan.productID) ?? Int.max
        }

        guard let current else { return candidate }
        return rank(for: candidate) < rank(for: current) ? candidate : current
    }

    private func activePlanFromLatestTransactions() async -> PaywallPlan? {
        var resolvedPlan: PaywallPlan?
        let productIDs = SubscriptionProductID.productIDs(from: PaywallConfigService.shared.config)

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

        return resolvedPlan
    }

    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.isLifetime else { continue }
            resolvedPlan = preferredPlan(resolvedPlan, plan)
        }

        return resolvedPlan
    }

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

            if isRevokedAppTransaction(transaction) {
                revokePremiumAccess()
                await transaction.finish()
                await refreshPremiumAccess()
                continue
            }

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

            await transaction.finish()

            if trustedPlan != nil {
                await reconcilePremiumAccess(trustedPlan: trustedPlan)
            } else {
                await refreshPremiumAccess()
            }
        }
    }

    private func isRevokedAppTransaction(_ transaction: StoreKit.Transaction) -> Bool {
        let knownProductIDs = Set(SubscriptionProductID.productIDs(from: PaywallConfigService.shared.config))
        return knownProductIDs.contains(transaction.productID) && transaction.revocationDate != nil
    }
}