macos-ai-menubar-usage/AIMenubarUsage/Services/UsageRefreshService.swift

import Foundation
import Combine

@MainActor
final class UsageRefreshService: ObservableObject {
    static let shared = UsageRefreshService()

    @Published private(set) var providerStates: [ProviderID: ProviderDisplayState] = [:]
    @Published private(set) var lastRefreshAt: Date?
    @Published private(set) var isRefreshing = false

    private let settings = AppSettings.shared
    private let cache = UsageCache.shared
    private let providers: [ProviderID: any UsageProvider] = [
        .cursor: CursorProvider(),
        .claude: ClaudeProvider(),
    ]
    private var refreshTask: Task<Void, Never>?
    private var timerCancellable: AnyCancellable?

    private init() {
        startTimer()
    }

    func startTimer() {
        timerCancellable?.cancel()
        timerCancellable = Timer.publish(every: TimeInterval(settings.refreshInterval.rawValue), on: .main, in: .common)
            .autoconnect()
            .sink { [weak self] _ in
                Task { await self?.refresh() }
            }
    }

    func restartTimer() {
        startTimer()
    }

    func refresh() async {
        guard !isRefreshing else { return }
        isRefreshing = true

        let enabled = settings.enabledProviders
        for provider in enabled {
            providerStates[provider] = .loading(provider)
        }

        await withTaskGroup(of: (ProviderID, ProviderDisplayState).self) { group in
            for providerID in enabled {
                group.addTask {
                    await self.fetchProvider(providerID)
                }
            }

            for await (providerID, state) in group {
                providerStates[providerID] = state
            }
        }

        for provider in ProviderID.allCases where !enabled.contains(provider) {
            providerStates.removeValue(forKey: provider)
        }

        lastRefreshAt = Date()
        isRefreshing = false
    }

    private func fetchProvider(_ providerID: ProviderID) async -> (ProviderID, ProviderDisplayState) {
        guard let provider = providers[providerID] else {
            return (providerID, ProviderDisplayState.failure(providerID, message: "Unknown provider"))
        }

        let providerSettings = settings.settings(for: providerID)

        do {
            let credentials = try provider.resolveCredentials(authMode: providerSettings.authMode)
            let snapshot = try await provider.fetchUsage(credentials: credentials, cache: cache)
            return (providerID, .success(snapshot))
        } catch {
            return (providerID, ProviderDisplayState.failure(providerID, message: error.localizedDescription))
        }
    }
}