FloatingBarUsageLimiter.swift

import Foundation

/// Tracks AI chat/query usage and enforces limits using the server-side quota.
///
/// Shared between the floating bar (Ask omi / PTT queries) and the main chat page
/// (ChatProvider.sendMessage). The server endpoint `/v1/users/me/usage-quota` is
/// the single source of truth for the current month's usage and plan limit.
@MainActor
final class FloatingBarUsageLimiter: ObservableObject {
  static let shared = FloatingBarUsageLimiter()

  /// Local mirror for avoiding needless requests after the device budget is
  /// exhausted. The backend independently resolves the subscription and is the
  /// spend authority; this cached value can only make the client more permissive.
  static func proactiveBudgetMultiplier(defaults: UserDefaults = .standard, now: Date = Date()) -> Int {
    switch defaults.string(forKey: .floatingBarCachedPlan) {
    case SubscriptionPlanType.architect.rawValue, SubscriptionPlanType.pro.rawValue:
      return 4
    case SubscriptionPlanType.operator.rawValue:
      return 2
    case SubscriptionPlanType.unlimited.rawValue:
      let grandfatherUntil = defaults.double(forKey: .floatingBarCachedDesktopGrandfatherUntil)
      return grandfatherUntil > now.timeIntervalSince1970 ? 2 : 1
    default:
      return 1
    }
  }

  @Published private(set) var hasPaidPlan: Bool = false

  /// Server-reported quota snapshot, plus an optimistic local delta for queries
  /// sent since the last server sync.
  @Published private(set) var serverQuota: APIClient.ChatUsageQuota?
  @Published private(set) var optimisticDelta: Int = 0

  init() {
    hasPaidPlan =
      UserDefaults.standard.string(forKey: .floatingBarCachedPlan)
      .map { SubscriptionPlanType(rawValue: $0).hasPaidCapability } ?? false
  }

  /// Fetch the user's subscription plan and usage quota from the backend.
  /// Call on app launch, sign-in, and after checkout completes.
  func fetchPlan() async {
    do {
      let response = try await APIClient.shared.getUserSubscription()
      applyPlan(
        plan: response.subscription.plan,
        status: response.subscription.status,
        desktopGrandfatherUntil: response.desktopGrandfatherUntil)
    } catch {
      log("FloatingBarUsageLimiter: failed to fetch plan: \(error.localizedDescription)")
    }
    await syncQuota()
  }

  /// Sync quota from the server, resetting the optimistic delta.
  func syncQuota() async {
    if let quota = await APIClient.shared.fetchChatUsageQuota() {
      applyQuota(quota)
    }
  }

  /// Apply a quota snapshot directly (used by syncQuota and tests).
  func applyQuota(_ quota: APIClient.ChatUsageQuota) {
    serverQuota = quota
    optimisticDelta = 0
  }

  /// Update cached plan directly from an already-fetched subscription (no extra API call).
  func applyPlan(
    plan: SubscriptionPlanType,
    status: SubscriptionStatusType,
    desktopGrandfatherUntil: Int? = nil
  ) {
    hasPaidPlan = plan.hasPaidCapability && status == .active
    // A verified active subscription is authoritative over a stale
    // trial/usage flag. Neo uses the Free Desktop floor for non-premium
    // features, but it is still paid and must never remain blocked from
    // audio capture while the app waits for another trial-metadata poll.
    if hasPaidPlan {
      AppState.current?.isPaywalled = false
      UserDefaults.standard.set(false, forKey: .desktopIsPaywalled)
    }
    if hasPaidPlan, serverQuota?.planType == SubscriptionPlanType.basic.rawValue {
      serverQuota = nil
      optimisticDelta = 0
    }
    // Persist only an active entitlement. Caching an inactive paid plan would
    // incorrectly restore both paid access and the larger proactive budget on
    // the next launch before the subscription refresh completes.
    let shouldPreservePlanIdentity: Bool
    if case .unknown = plan {
      shouldPreservePlanIdentity = true
    } else {
      shouldPreservePlanIdentity = hasPaidPlan
    }
    UserDefaults.standard.set(
      shouldPreservePlanIdentity ? plan.rawValue : SubscriptionPlanType.basic.rawValue,
      forKey: .floatingBarCachedPlan)
    if hasPaidPlan, plan == .unlimited, let desktopGrandfatherUntil {
      UserDefaults.standard.set(
        desktopGrandfatherUntil,
        forKey: .floatingBarCachedDesktopGrandfatherUntil)
    } else {
      UserDefaults.standard.removeObject(forKey: .floatingBarCachedDesktopGrandfatherUntil)
    }
  }

  /// Reset all quota state on sign-out so the next user starts clean.
  func reset() {
    serverQuota = nil
    optimisticDelta = 0
    hasPaidPlan = false
    UserDefaults.standard.removeObject(forKey: .floatingBarCachedPlan)
    UserDefaults.standard.removeObject(forKey: .floatingBarCachedDesktopGrandfatherUntil)
  }

  var isLimitReached: Bool {
    // BYOK users pay their own LLM bill and are never limited. Honor local
    // BYOK state so a heartbeat-lagged server quota (allowed=false right
    // after activation) can't block chat for a fully-configured BYOK user.
    if APIKeyService.isByokActive {
      return false
    }
    guard let quota = serverQuota else {
      // No server data yet — allow the query (server will enforce).
      return false
    }
    // An overage plan bills the excess instead of refusing it: the server serves
    // the request and `enforce_chat_quota` never raises, so blocking here would
    // deny a send the backend would have answered.
    if quota.isOveragePlan == true {
      return false
    }
    if quota.allowed {
      // Optimistic delta only applies to question-based quotas.
      // For cost_usd (Architect/Pro), we can't estimate cost per query
      // locally — rely on the server snapshot alone.
      guard quota.unit == "questions", let limit = quota.limit else { return false }
      return (quota.used + Double(optimisticDelta)) >= limit
    }
    return true
  }

  var remainingQueries: Int {
    guard let quota = serverQuota else { return .max }
    guard quota.unit == "questions", let limit = quota.limit else { return .max }
    return max(0, Int(limit - quota.used) - optimisticDelta)
  }

  /// Human-readable limit text for error messages.
  var limitDescription: String {
    guard let quota = serverQuota, let limit = quota.limit else {
      return "your monthly free message limit"
    }
    if quota.unit == "cost_usd" {
      return String(format: "your $%.0f %@ monthly spend limit", limit, quota.plan)
    }
    return "\(Int(limit)) \(quota.plan) messages this month"
  }

  /// Record a query. Call after successfully sending a query from the floating bar
  /// OR the main chat page — both surfaces share this pool.
  func recordQuery() {
    optimisticDelta += 1
  }
}