nucleic/Sources/NucleicCore/SubscriptionUsage.swift

import Foundation

#if canImport(Security)
import Security
#endif
#if canImport(FoundationNetworking)
import FoundationNetworking
#endif

/// One subscription rate-limit window from the usage endpoint.
public struct UsageWindow: Sendable, Equatable {
    /// Percent of the window consumed, 0–100.
    public let utilization: Double
    /// When this window rolls over and frees capacity (nil if the server omits it).
    public let resetsAt: Date?

    public init(utilization: Double, resetsAt: Date?) {
        self.utilization = utilization
        self.resetsAt = resetsAt
    }

    /// The utilization to show at `now`, correcting for stale snapshots across a reset.
    /// Usage is polled coarsely (every ~2 min; see `AppStore.startQuotaPolling`), so once
    /// `resetsAt` has passed the cached snapshot still reports the pre-reset percentage
    /// until the next fetch lands. The window has already rolled over and freed its
    /// capacity by then, so a lapsed window reads as empty (0%) rather than showing the
    /// stale, now-wrong value. Windows with no `resetsAt` are reported as-is.
    public func utilization(at now: Date) -> Double {
        if let resetsAt, now >= resetsAt { return 0 }
        return utilization
    }
}

/// Quantitative Claude **subscription** usage (Pro/Max) — the same numbers the
/// `claude` CLI's `/usage` screen shows: how much of your rolling 5-hour and weekly
/// windows you've consumed. This is plan usage, *not* API billing.
///
/// Sourced from the undocumented `GET /api/oauth/usage` endpoint, authenticated with the
/// subscription OAuth token from Nucleic's mediated credential store (`ClaudeLoginKeychain`) —
/// refreshed on demand by `ClaudeCredentialBroker.freshAccessToken()`. The endpoint is internal to
/// Claude Code and may change without notice, so every caller must tolerate
/// `SubscriptionUsageFetcher.fetch()` throwing and fall back to the coarser `rate_limit_event`
/// status (see `AppStore.latestRateLimit`).
public struct SubscriptionUsage: Sendable, Equatable {
    /// The rolling 5-hour ("session") limit.
    public let fiveHour: UsageWindow?
    /// The rolling 7-day ("weekly") limit, across all models.
    public let sevenDay: UsageWindow?
    /// The weekly Opus sub-limit (nil when the plan has no separate Opus cap).
    public let sevenDayOpus: UsageWindow?
    /// The weekly Sonnet sub-limit.
    public let sevenDaySonnet: UsageWindow?

    public init(
        fiveHour: UsageWindow?, sevenDay: UsageWindow?,
        sevenDayOpus: UsageWindow?, sevenDaySonnet: UsageWindow?
    ) {
        self.fiveHour = fiveHour
        self.sevenDay = sevenDay
        self.sevenDayOpus = sevenDayOpus
        self.sevenDaySonnet = sevenDaySonnet
    }

    /// The most-constrained window's utilization — what a single headline indicator
    /// should track, since hitting *any* window throttles you.
    public var peakUtilization: Double? {
        [fiveHour, sevenDay, sevenDayOpus, sevenDaySonnet]
            .compactMap { $0?.utilization }.max()
    }
}

public enum SubscriptionUsageError: Error, Sendable {
    /// No Claude Code OAuth token in the keychain (logged out, or using an API key).
    case noToken
    case badResponse
    case http(Int)
}

/// Fetches `SubscriptionUsage` from Anthropic's OAuth usage endpoint. `nonisolated`
/// and stateless: safe to call from the main actor — the network await runs off-main.
public enum SubscriptionUsageFetcher {
    private static let endpoint = URL(string: "https://api.anthropic.com/api/oauth/usage")!
    /// The OAuth beta the CLI sends on this endpoint (mirrors claude 2.1.x).
    private static let oauthBeta = "oauth-2025-04-20"

    /// Fetch using whatever token is currently in the mediated store. Prefer ``fetch(token:)`` with
    /// a broker-refreshed token so an expired access token self-heals instead of silently failing.
    public static func fetch() async throws -> SubscriptionUsage {
        guard let token = oauthToken() else { throw SubscriptionUsageError.noToken }
        return try await fetch(token: token)
    }

    /// Fetch with an explicit access token — the one ``ClaudeCredentialBroker/freshAccessToken()``
    /// returns after refreshing through Nucleic's own OAuth. Keeps usage polling self-sufficient:
    /// it no longer needs a session turn (or the external CLI) to rotate the keychain token first.
    public static func fetch(token: String) async throws -> SubscriptionUsage {
        var request = URLRequest(url: endpoint)
        request.timeoutInterval = 10
        request.httpMethod = "GET"
        request.setValue("Bearer \(token)", forHTTPHeaderField: "Authorization")
        request.setValue("application/json", forHTTPHeaderField: "Content-Type")
        request.setValue(oauthBeta, forHTTPHeaderField: "anthropic-beta")
        request.setValue("2023-06-01", forHTTPHeaderField: "anthropic-version")

        let (data, response) = try await URLSession.shared.data(for: request)
        guard let http = response as? HTTPURLResponse else { throw SubscriptionUsageError.badResponse }
        guard http.statusCode == 200 else { throw SubscriptionUsageError.http(http.statusCode) }
        return try parse(data)
    }

    /// Reads the access token from Nucleic's prompt-free credential store. On Linux the same
    /// abstraction is backed by Claude's runner-side credentials file.
    static func oauthToken() -> String? {
        guard let json = ClaudeLoginKeychain.read(),
              let data = json.data(using: .utf8),
              let obj = try? JSONSerialization.jsonObject(with: data) as? [String: Any],
              let oauth = obj["claudeAiOauth"] as? [String: Any],
              let token = oauth["accessToken"] as? String
        else { return nil }
        return token
    }

    static func parse(_ data: Data) throws -> SubscriptionUsage {
        guard let json = try JSONSerialization.jsonObject(with: data) as? [String: Any] else {
            throw SubscriptionUsageError.badResponse
        }
        func window(_ key: String) -> UsageWindow? {
            guard let obj = json[key] as? [String: Any],
                  let util = obj["utilization"] as? Double else { return nil }
            let resets = (obj["resets_at"] as? String).flatMap(parseDate)
            return UsageWindow(utilization: util, resetsAt: resets)
        }
        return SubscriptionUsage(
            fiveHour: window("five_hour"),
            sevenDay: window("seven_day"),
            sevenDayOpus: window("seven_day_opus"),
            sevenDaySonnet: window("seven_day_sonnet"))
    }

    /// The endpoint stamps microsecond fractional seconds (`…851919+00:00`); try the
    /// fractional-seconds parser first, then fall back to whole seconds.
    private static func parseDate(_ string: String) -> Date? {
        let withFraction = ISO8601DateFormatter()
        withFraction.formatOptions = [.withInternetDateTime, .withFractionalSeconds]
        if let date = withFraction.date(from: string) { return date }
        let plain = ISO8601DateFormatter()
        plain.formatOptions = [.withInternetDateTime]
        return plain.date(from: string)
    }
}