macos-ai-menubar-usage/AIMenubarUsage/Providers/ClaudeProvider.swift

import Foundation

enum ClaudeAPIError: Error, LocalizedError {
    case notAuthenticated
    case invalidResponse
    case rateLimited
    case httpError(statusCode: Int)
    case cliUnavailable

    var errorDescription: String? {
        switch self {
        case .notAuthenticated:
            "Not authenticated. Run `claude auth login` or add credentials in Settings."
        case .invalidResponse:
            "Could not parse usage data from Claude."
        case .rateLimited:
            "Claude usage API rate limited. Showing cached data if available."
        case .httpError(let statusCode):
            "Claude API returned HTTP \(statusCode)."
        case .cliUnavailable:
            "Claude CLI usage is unavailable."
        }
    }
}

struct ClaudeUsageWindow: Codable, Sendable {
    let utilization: Double?
    let resetsAt: String?

    enum CodingKeys: String, CodingKey {
        case utilization
        case resetsAt = "resets_at"
    }
}

struct ClaudeOAuthUsageResponse: Codable, Sendable {
    let fiveHour: ClaudeUsageWindow?
    let sevenDay: ClaudeUsageWindow?
    let extraUsage: ClaudeUsageWindow?
    let sevenDayOpus: ClaudeUsageWindow?
    let sevenDaySonnet: ClaudeUsageWindow?

    enum CodingKeys: String, CodingKey {
        case fiveHour = "five_hour"
        case sevenDay = "seven_day"
        case extraUsage = "extra_usage"
        case sevenDayOpus = "seven_day_opus"
        case sevenDaySonnet = "seven_day_sonnet"
    }
}

struct ClaudeUsageReportBucket: Decodable, Sendable {
    let inputTokens: Int?
    let outputTokens: Int?

    enum CodingKeys: String, CodingKey {
        case inputTokens = "input_tokens"
        case outputTokens = "output_tokens"
    }
}

struct ClaudeUsageReportData: Decodable, Sendable {
    let data: [ClaudeUsageReportBucket]?
}

struct ClaudeUsageReportResponse: Decodable, Sendable {
    let data: [ClaudeUsageReportData]?
}

struct ClaudeCostBucket: Decodable, Sendable {
    let amount: String?
    let currency: String?
}

struct ClaudeCostReportData: Decodable, Sendable {
    let data: [ClaudeCostBucket]?
}

struct ClaudeCostReportResponse: Decodable, Sendable {
    let data: [ClaudeCostReportData]?
}

enum ClaudeAPI {
    private static let oauthUsageURL = URL(string: "https://api.anthropic.com/api/oauth/usage")!
    private static let usageReportURL = URL(string: "https://api.anthropic.com/v1/organizations/usage_report/messages")!
    private static let costReportURL = URL(string: "https://api.anthropic.com/v1/organizations/cost_report")!

    static func fetchOAuthUsage(accessToken: String) async throws -> ClaudeOAuthUsageResponse {
        var request = URLRequest(url: oauthUsageURL)
        request.httpMethod = "GET"
        request.setValue("Bearer \(accessToken)", forHTTPHeaderField: "Authorization")
        request.setValue("application/json", forHTTPHeaderField: "Accept")
        request.setValue("oauth-2025-04-20", forHTTPHeaderField: "anthropic-beta")
        request.setValue("claude-cli/2.1.119 (external, cli)", forHTTPHeaderField: "User-Agent")

        let (data, response) = try await URLSession.shared.data(for: request)
        guard let httpResponse = response as? HTTPURLResponse else {
            throw ClaudeAPIError.invalidResponse
        }

        switch httpResponse.statusCode {
        case 200:
            break
        case 401:
            throw ClaudeAPIError.notAuthenticated
        case 429:
            throw ClaudeAPIError.rateLimited
        default:
            throw ClaudeAPIError.httpError(statusCode: httpResponse.statusCode)
        }

        do {
            return try JSONDecoder().decode(ClaudeOAuthUsageResponse.self, from: data)
        } catch {
            throw ClaudeAPIError.invalidResponse
        }
    }

    static func fetchAdminUsage(apiKey: String) async throws -> (inputTokens: Int, outputTokens: Int, costUSD: Double) {
        let ending = Date()
        let starting = Calendar.current.date(byAdding: .day, value: -7, to: ending) ?? ending
        let formatter = ISO8601DateFormatter()

        var usageComponents = URLComponents(url: usageReportURL, resolvingAgainstBaseURL: false)!
        usageComponents.queryItems = [
            URLQueryItem(name: "starting_at", value: formatter.string(from: starting)),
            URLQueryItem(name: "ending_at", value: formatter.string(from: ending)),
            URLQueryItem(name: "bucket_width", value: "1d"),
        ]

        var costComponents = URLComponents(url: costReportURL, resolvingAgainstBaseURL: false)!
        costComponents.queryItems = [
            URLQueryItem(name: "starting_at", value: formatter.string(from: starting)),
            URLQueryItem(name: "ending_at", value: formatter.string(from: ending)),
            URLQueryItem(name: "bucket_width", value: "1d"),
        ]

        async let usageData = adminRequest(url: usageComponents.url!, apiKey: apiKey)
        async let costData = adminRequest(url: costComponents.url!, apiKey: apiKey)

        let (usageRaw, costRaw) = try await (usageData, costData)
        let usage = try JSONDecoder().decode(ClaudeUsageReportResponse.self, from: usageRaw)
        let cost = try JSONDecoder().decode(ClaudeCostReportResponse.self, from: costRaw)

        var inputTokens = 0
        var outputTokens = 0
        for group in usage.data ?? [] {
            for bucket in group.data ?? [] {
                inputTokens += bucket.inputTokens ?? 0
                outputTokens += bucket.outputTokens ?? 0
            }
        }

        var costUSD = 0.0
        for group in cost.data ?? [] {
            for bucket in group.data ?? [] {
                if let amount = bucket.amount, let value = Double(amount) {
                    costUSD += value
                }
            }
        }

        return (inputTokens, outputTokens, costUSD)
    }

    private static func adminRequest(url: URL, apiKey: String) async throws -> Data {
        var request = URLRequest(url: url)
        request.httpMethod = "GET"
        request.setValue(apiKey, forHTTPHeaderField: "x-api-key")
        request.setValue("2023-06-01", forHTTPHeaderField: "anthropic-version")

        let (data, response) = try await URLSession.shared.data(for: request)
        guard let httpResponse = response as? HTTPURLResponse else {
            throw ClaudeAPIError.invalidResponse
        }

        switch httpResponse.statusCode {
        case 200:
            return data
        case 401, 403:
            throw ClaudeAPIError.notAuthenticated
        case 429:
            throw ClaudeAPIError.rateLimited
        default:
            throw ClaudeAPIError.httpError(statusCode: httpResponse.statusCode)
        }
    }

    static func snapshot(from usage: ClaudeOAuthUsageResponse, fetchedAt: Date, isStale: Bool) -> ProviderUsageSnapshot {
        var values: [UsageMetric: MetricValue] = [:]
        let isoFormatter = ISO8601DateFormatter()
        isoFormatter.formatOptions = [.withInternetDateTime, .withFractionalSeconds]

        if let window = usage.fiveHour, let utilization = window.utilization {
            values[.claudeFiveHourUtilization] = .percent(utilization * 100)
        }
        if let reset = usage.fiveHour?.resetsAt,
           let date = isoFormatter.date(from: reset) ?? ISO8601DateFormatter().date(from: reset) {
            values[.claudeFiveHourResetsAt] = .date(date)
        }
        if let window = usage.sevenDay, let utilization = window.utilization {
            values[.claudeSevenDayUtilization] = .percent(utilization * 100)
        }
        if let reset = usage.sevenDay?.resetsAt,
           let date = isoFormatter.date(from: reset) ?? ISO8601DateFormatter().date(from: reset) {
            values[.claudeSevenDayResetsAt] = .date(date)
        }
        if let window = usage.extraUsage, let utilization = window.utilization {
            values[.claudeExtraUsage] = .percent(utilization * 100)
        }
        if let window = usage.sevenDayOpus, let utilization = window.utilization {
            values[.claudeSevenDayOpusUtilization] = .percent(utilization * 100)
        }
        if let window = usage.sevenDaySonnet, let utilization = window.utilization {
            values[.claudeSevenDaySonnetUtilization] = .percent(utilization * 100)
        }

        return ProviderUsageSnapshot(
            providerID: .claude,
            fetchedAt: fetchedAt,
            values: values,
            isStale: isStale
        )
    }

    static func adminSnapshot(inputTokens: Int, outputTokens: Int, costUSD: Double, fetchedAt: Date, isStale: Bool) -> ProviderUsageSnapshot {
        ProviderUsageSnapshot(
            providerID: .claude,
            fetchedAt: fetchedAt,
            values: [
                .claudeAdminInputTokens: .integer(inputTokens),
                .claudeAdminOutputTokens: .integer(outputTokens),
                .claudeAdminCostUSD: .currency(costUSD),
            ],
            isStale: isStale
        )
    }
}

struct ClaudeProvider: UsageProvider {
    let id: ProviderID = .claude
    let displayName = "Claude"

    func resolveCredentials(authMode: AuthMode) throws -> AuthCredentials {
        switch authMode {
        case .auto:
            return .claudeCLI
        case .manualSession:
            let account = KeychainService.keychainAccount(provider: .claude, authMode: .manualSession)
            guard let token = KeychainService.load(account: account), !token.isEmpty else {
                throw ClaudeAPIError.notAuthenticated
            }
            return .claudeOAuth(accessToken: token)
        case .adminAPI:
            let account = KeychainService.keychainAccount(provider: .claude, authMode: .adminAPI)
            guard let apiKey = KeychainService.load(account: account), !apiKey.isEmpty else {
                throw ClaudeAPIError.notAuthenticated
            }
            return .claudeAdminAPI(apiKey: apiKey)
        }
    }

    func fetchUsage(credentials: AuthCredentials, cache: UsageCache) async throws -> ProviderUsageSnapshot {
        switch credentials {
        case .claudeCLI:
            return try await fetchCLIUsage(cache: cache)
        case .claudeOAuth(let accessToken):
            return try await fetchOAuthUsage(accessToken: accessToken, cache: cache)
        case .claudeAdminAPI(let apiKey):
            return try await fetchAdminUsage(apiKey: apiKey, cache: cache)
        default:
            throw ClaudeAPIError.invalidResponse
        }
    }

    private func fetchCLIUsage(cache: UsageCache) async throws -> ProviderUsageSnapshot {
        if let cached = await cache.load(provider: .claude, maxAge: 300),
           let parsed = try? JSONDecoder().decode(ClaudeCLIParsedUsage.self, from: cached.payload) {
            return ClaudeCLIUsageService.snapshot(from: parsed, fetchedAt: cached.fetchedAt, isStale: false)
        }

        do {
            let snapshot = try await ClaudeCLIUsageService.fetchUsage()
            if let parsed = cliParsedUsage(from: snapshot),
               let data = try? JSONEncoder().encode(parsed) {
                await cache.save(provider: .claude, payload: data)
            }
            return snapshot
        } catch {
            if let stale = await cache.loadStale(provider: .claude),
               let parsed = try? JSONDecoder().decode(ClaudeCLIParsedUsage.self, from: stale.payload) {
                return ClaudeCLIUsageService.snapshot(from: parsed, fetchedAt: stale.fetchedAt, isStale: true)
            }
            throw error
        }
    }

    private func cliParsedUsage(from snapshot: ProviderUsageSnapshot) -> ClaudeCLIParsedUsage? {
        func percent(_ metric: UsageMetric) -> Double? {
            if case .percent(let value)? = snapshot.value(for: metric) {
                return value
            }
            return nil
        }
        func text(_ metric: UsageMetric) -> String? {
            if case .text(let value)? = snapshot.value(for: metric) {
                return value
            }
            return nil
        }

        var modelWeekly: [String: Double] = [:]
        if let opus = percent(.claudeSevenDayOpusUtilization) {
            modelWeekly["Opus"] = opus
        }
        if let sonnet = percent(.claudeSevenDaySonnetUtilization) {
            modelWeekly["Sonnet"] = sonnet
        }

        return ClaudeCLIParsedUsage(
            sessionPercent: percent(.claudeFiveHourUtilization),
            weeklyPercent: percent(.claudeSevenDayUtilization),
            weeklyResetText: text(.claudeSevenDayResetsAt),
            modelWeekly: modelWeekly
        )
    }

    private func fetchOAuthUsage(accessToken: String, cache: UsageCache) async throws -> ProviderUsageSnapshot {
        if let cached = await cache.load(provider: .claude, maxAge: 300),
           let usage = try? JSONDecoder().decode(ClaudeOAuthUsageResponse.self, from: cached.payload) {
            return ClaudeAPI.snapshot(from: usage, fetchedAt: cached.fetchedAt, isStale: false)
        }

        do {
            let usage = try await ClaudeAPI.fetchOAuthUsage(accessToken: accessToken)
            if let data = try? JSONEncoder().encode(usage) {
                await cache.save(provider: .claude, payload: data)
            }
            return ClaudeAPI.snapshot(from: usage, fetchedAt: Date(), isStale: false)
        } catch {
            if let stale = await cache.loadStale(provider: .claude),
               let usage = try? JSONDecoder().decode(ClaudeOAuthUsageResponse.self, from: stale.payload) {
                return ClaudeAPI.snapshot(from: usage, fetchedAt: stale.fetchedAt, isStale: true)
            }
            throw error
        }
    }

    private func fetchAdminUsage(apiKey: String, cache: UsageCache) async throws -> ProviderUsageSnapshot {
        struct AdminPayload: Codable {
            let inputTokens: Int
            let outputTokens: Int
            let costUSD: Double
        }

        if let cached = await cache.load(provider: .claude, maxAge: 300),
           let payload = try? JSONDecoder().decode(AdminPayload.self, from: cached.payload) {
            return ClaudeAPI.adminSnapshot(
                inputTokens: payload.inputTokens,
                outputTokens: payload.outputTokens,
                costUSD: payload.costUSD,
                fetchedAt: cached.fetchedAt,
                isStale: false
            )
        }

        do {
            let result = try await ClaudeAPI.fetchAdminUsage(apiKey: apiKey)
            let payload = AdminPayload(
                inputTokens: result.inputTokens,
                outputTokens: result.outputTokens,
                costUSD: result.costUSD
            )
            if let data = try? JSONEncoder().encode(payload) {
                await cache.save(provider: .claude, payload: data)
            }
            return ClaudeAPI.adminSnapshot(
                inputTokens: result.inputTokens,
                outputTokens: result.outputTokens,
                costUSD: result.costUSD,
                fetchedAt: Date(),
                isStale: false
            )
        } catch {
            if let stale = await cache.loadStale(provider: .claude),
               let payload = try? JSONDecoder().decode(AdminPayload.self, from: stale.payload) {
                return ClaudeAPI.adminSnapshot(
                    inputTokens: payload.inputTokens,
                    outputTokens: payload.outputTokens,
                    costUSD: payload.costUSD,
                    fetchedAt: stale.fetchedAt,
                    isStale: true
                )
            }
            throw error
        }
    }
}