steipete--codexbar/Sources/CodexBarCore/Providers/Grok/GrokProviderDescriptor.swift

import Foundation

public enum GrokProviderDescriptor {
    public static let descriptor: ProviderDescriptor = Self.makeDescriptor()

    static func makeDescriptor() -> ProviderDescriptor {
        ProviderDescriptor(
            id: .grok,
            metadata: ProviderMetadata(
                id: .grok,
                displayName: "Grok",
                sessionLabel: "Credits",
                weeklyLabel: "On-demand",
                opusLabel: nil,
                supportsOpus: false,
                supportsCredits: false,
                creditsHint: "",
                toggleTitle: "Show Grok usage",
                cliName: "grok",
                defaultEnabled: false,
                isPrimaryProvider: false,
                usesAccountFallback: false,
                browserCookieOrder: ProviderBrowserCookieDefaults.grokCookieImportOrder,
                dashboardURL: "https://grok.com/?_s=usage",
                changelogURL: "https://x.ai/news",
                statusPageURL: nil,
                statusLinkURL: "https://status.x.ai"),
            branding: ProviderBranding(
                iconStyle: .grok,
                iconResourceName: "ProviderIcon-grok",
                color: ProviderColor(red: 16 / 255, green: 163 / 255, blue: 127 / 255)),
            tokenCost: ProviderTokenCostConfig(
                supportsTokenCost: false,
                noDataMessage: { "Grok cost summary is not supported yet." }),
            fetchPlan: ProviderFetchPlan(
                sourceModes: [.auto, .cli, .web],
                pipeline: ProviderFetchPipeline(resolveStrategies: self.resolveStrategies)),
            cli: ProviderCLIConfig(
                name: "grok",
                versionDetector: { _ in GrokStatusProbe.detectVersion() }))
    }

    private static func resolveStrategies(context: ProviderFetchContext) async -> [any ProviderFetchStrategy] {
        switch context.sourceMode {
        case .auto:
            [GrokCLIFetchStrategy(), GrokWebFetchStrategy()]
        case .cli:
            [GrokCLIFetchStrategy()]
        case .web:
            [GrokWebFetchStrategy()]
        case .api, .oauth:
            []
        }
    }

    /// Returns a contextual label for Grok's primary usage bar ("Weekly" or "Monthly").
    /// Prefer the billing period duration when available; fall back to reset distance for
    /// web billing payloads that expose only a reset timestamp.
    public static func primaryLabel(window: RateWindow?, now: Date = .now) -> String? {
        if let minutes = window?.windowMinutes {
            return self.primaryLabel(duration: TimeInterval(minutes) * 60)
        }
        return self.primaryLabel(resetsAt: window?.resetsAt, now: now)
    }

    public static func primaryLabel(resetsAt: Date?, now: Date = .now) -> String? {
        guard let resetsAt else { return nil }
        return self.primaryLabel(duration: resetsAt.timeIntervalSince(now))
    }

    private static func primaryLabel(duration seconds: TimeInterval) -> String? {
        guard seconds > 3600 else { return nil }
        let days = Int((seconds / 86400).rounded(.toNearestOrAwayFromZero))
        if (4...12).contains(days) { return "Weekly" }
        if (20...45).contains(days) { return "Monthly" }
        return nil
    }
}

struct GrokCLIFetchStrategy: ProviderFetchStrategy {
    let id: String = "grok.cli"
    let kind: ProviderFetchKind = .cli

    func isAvailable(_ context: ProviderFetchContext) async -> Bool {
        BinaryLocator.resolveGrokBinary(env: context.env) != nil
    }

    func fetch(_ context: ProviderFetchContext) async throws -> ProviderFetchResult {
        let probe = GrokStatusProbe()
        let snap = try await probe.fetch(env: context.env)
        return self.makeResult(
            usage: snap.toUsageSnapshot(),
            sourceLabel: "grok-cli")
    }

    func shouldFallback(on _: Error, context: ProviderFetchContext) -> Bool {
        context.sourceMode == .auto
    }
}

struct GrokWebFetchStrategy: ProviderFetchStrategy {
    let id: String = "grok.web"
    let kind: ProviderFetchKind = .web

    static func canImportBrowserCookies(runtime: ProviderRuntime, env: [String: String]) -> Bool {
        runtime == .app || env["CODEXBAR_ALLOW_BROWSER_COOKIE_IMPORT"] == "1"
    }

    func isAvailable(_ context: ProviderFetchContext) async -> Bool {
        #if os(macOS)
        if Self.canImportBrowserCookies(runtime: context.runtime, env: context.env),
           GrokCookieImporter.hasSession(browserDetection: context.browserDetection)
        {
            return true
        }
        #endif
        return FileManager.default.fileExists(atPath: GrokCredentialsStore.authFileURL(env: context.env).path)
    }

    func fetch(_ context: ProviderFetchContext) async throws -> ProviderFetchResult {
        let (webBilling, sourceLabel, authenticatedByAuthFile) = try await self.fetchWebBilling(context: context)
        let credentials = Self.credentialsForWebBillingSnapshot(
            credentials: try? GrokCredentialsStore.load(env: context.env),
            authenticatedByAuthFile: authenticatedByAuthFile)
        let snapshot = GrokUsageSnapshot(
            billing: nil,
            webBilling: webBilling,
            credentials: GrokStatusProbe.credentialsForSnapshot(
                credentials: credentials,
                billing: nil,
                webBilling: webBilling),
            localSummary: GrokLocalSessionScanner.summarize(env: context.env),
            cliVersion: GrokStatusProbe.detectVersion(env: context.env),
            updatedAt: Date())
        return self.makeResult(
            usage: snapshot.toUsageSnapshot(),
            sourceLabel: sourceLabel)
    }

    private func fetchWebBilling(context: ProviderFetchContext) async throws -> (
        snapshot: GrokWebBillingSnapshot,
        sourceLabel: String,
        authenticatedByAuthFile: Bool)
    {
        let credentialsResult: Result<GrokCredentials, Error> = Result {
            try GrokCredentialsStore.load(env: context.env)
        }
        let browserCredentials = try? credentialsResult.get()

        #if os(macOS)
        if Self.canImportBrowserCookies(runtime: context.runtime, env: context.env) {
            var lastCookieError: Error?
            do {
                let sessions = try GrokCookieImporter.importSessions(browserDetection: context.browserDetection)
                let (snapshot, sourceLabel) = try await Self.fetchFirstValidCookieSession(
                    sessions,
                    credentials: browserCredentials)
                return (snapshot, sourceLabel, false)
            } catch {
                lastCookieError = error
            }
            if browserCredentials == nil {
                if FileManager.default.fileExists(
                    atPath: GrokCredentialsStore.authFileURL(env: context.env).path)
                {
                    _ = try credentialsResult.get()
                }
                throw lastCookieError ?? GrokWebBillingError.missingCredentials
            }
        }
        #endif

        let authCredentials = try credentialsResult.get()
        guard !authCredentials.isExpired else {
            throw GrokWebBillingError.missingCredentials
        }
        let snapshot = try await GrokWebBillingFetcher.fetch(credentials: authCredentials)
        return (snapshot, "grok-web", true)
    }

    static func credentialsForWebBillingSnapshot(
        credentials: GrokCredentials?,
        authenticatedByAuthFile: Bool) -> GrokCredentials?
    {
        authenticatedByAuthFile ? credentials : nil
    }

    #if os(macOS)
    static func fetchFirstValidCookieSession(
        _ sessions: [GrokCookieImporter.SessionInfo],
        credentials: GrokCredentials? = nil,
        fetch: ((String, GrokCredentials?) async throws -> GrokWebBillingSnapshot)? = nil) async throws
        -> (GrokWebBillingSnapshot, String)
    {
        let fetchSnapshot = fetch ?? { cookieHeader, credentials in
            try await GrokWebBillingFetcher.fetch(
                cookieHeader: cookieHeader,
                credentials: credentials)
        }
        var lastError: Error?
        for session in sessions {
            for authCredentials in Self.cookieAuthAttempts(credentials: credentials) {
                do {
                    let snapshot = try await fetchSnapshot(session.cookieHeader, authCredentials)
                    return (snapshot, session.sourceLabel)
                } catch {
                    lastError = error
                }
            }
        }
        throw lastError ?? GrokWebBillingError.missingCredentials
    }

    static func cookieAuthAttempts(credentials: GrokCredentials?) -> [GrokCredentials?] {
        guard let credentials, !credentials.isExpired else { return [nil] }
        return [credentials, nil]
    }
    #endif

    func shouldFallback(on _: Error, context _: ProviderFetchContext) -> Bool {
        false
    }
}