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

import Foundation

enum ClaudeCLIError: Error, LocalizedError {
    case cliNotFound
    case commandFailed(String)
    case invalidResponse
    case notAuthenticated

    var errorDescription: String? {
        switch self {
        case .cliNotFound:
            "Claude Code CLI not found. Install it from claude.ai/code or set CLAUDE_CLI_PATH."
        case .commandFailed(let details):
            "Claude CLI failed: \(details)"
        case .invalidResponse:
            "Could not parse usage output from Claude CLI."
        case .notAuthenticated:
            "Claude CLI is not signed in. Run `claude auth login` in Terminal."
        }
    }
}

struct ClaudeCLIPrintResponse: Decodable, Sendable {
    let isError: Bool?
    let result: String?
    let subtype: String?

    enum CodingKeys: String, CodingKey {
        case isError = "is_error"
        case result
        case subtype
    }
}

struct ClaudeCLIParsedUsage: Codable, Sendable {
    let sessionPercent: Double?
    let weeklyPercent: Double?
    let weeklyResetText: String?
    let modelWeekly: [String: Double]
}

enum ClaudeCLIUsageService {
    private static let timeoutSeconds: TimeInterval = 45

    static func resolveClaudeExecutable() -> String {
        if let configured = ProcessInfo.processInfo.environment["CLAUDE_CLI_PATH"],
           !configured.isEmpty,
           FileManager.default.isExecutableFile(atPath: configured) {
            return configured
        }

        let home = FileManager.default.homeDirectoryForCurrentUser.path
        let candidates = [
            "\(home)/.local/bin/claude",
            "/opt/homebrew/bin/claude",
            "/usr/local/bin/claude",
        ]

        for candidate in candidates where FileManager.default.isExecutableFile(atPath: candidate) {
            return candidate
        }

        return "claude"
    }

    static func fetchUsage() async throws -> ProviderUsageSnapshot {
        let output = try await runUsageCommand()
        let parsed = try parsePrintResponse(output)
        return snapshot(from: parsed, fetchedAt: Date(), isStale: false)
    }

    static func runUsageCommand() async throws -> String {
        let claudePath = resolveClaudeExecutable()
        let process = Process()
        process.executableURL = URL(fileURLWithPath: claudePath)
        process.arguments = ["-p", "/usage", "--output-format", "json"]

        var environment = ProcessInfo.processInfo.environment
        let home = FileManager.default.homeDirectoryForCurrentUser.path
        let extraPaths = [
            URL(fileURLWithPath: claudePath).deletingLastPathComponent().path,
            "\(home)/.local/bin",
            "/opt/homebrew/bin",
            "/usr/local/bin",
        ]
        let existingPath = environment["PATH"] ?? "/usr/bin:/bin"
        environment["PATH"] = (extraPaths + [existingPath]).joined(separator: ":")
        process.environment = environment

        let stdout = Pipe()
        let stderr = Pipe()
        process.standardOutput = stdout
        process.standardError = stderr

        try process.run()

        let deadline = Date().addingTimeInterval(timeoutSeconds)
        while process.isRunning, Date() < deadline {
            try await Task.sleep(nanoseconds: 100_000_000)
        }

        if process.isRunning {
            process.terminate()
            throw ClaudeCLIError.commandFailed("timed out after \(Int(timeoutSeconds))s")
        }

        let outputData = stdout.fileHandleForReading.readDataToEndOfFile()
        let errorData = stderr.fileHandleForReading.readDataToEndOfFile()
        let output = String(data: outputData, encoding: .utf8) ?? ""
        let errorOutput = String(data: errorData, encoding: .utf8) ?? ""

        guard process.terminationStatus == 0 else {
            let details = errorOutput.trimmingCharacters(in: .whitespacesAndNewlines)
            if claudePath == "claude", details.contains("not found") || details.contains("No such file") {
                throw ClaudeCLIError.cliNotFound
            }
            if details.localizedCaseInsensitiveContains("not logged in")
                || details.localizedCaseInsensitiveContains("not authenticated")
                || details.localizedCaseInsensitiveContains("login") {
                throw ClaudeCLIError.notAuthenticated
            }
            throw ClaudeCLIError.commandFailed(details.isEmpty ? "exit \(process.terminationStatus)" : details)
        }

        guard !output.trimmingCharacters(in: .whitespacesAndNewlines).isEmpty else {
            throw ClaudeCLIError.invalidResponse
        }

        return output
    }

    static func parsePrintResponse(_ output: String) throws -> ClaudeCLIParsedUsage {
        guard let data = output.data(using: .utf8) else {
            throw ClaudeCLIError.invalidResponse
        }

        let response = try JSONDecoder().decode(ClaudeCLIPrintResponse.self, from: data)
        if response.isError == true {
            throw ClaudeCLIError.commandFailed(response.result ?? "unknown CLI error")
        }

        guard let result = response.result, !result.isEmpty else {
            throw ClaudeCLIError.invalidResponse
        }

        if result.localizedCaseInsensitiveContains("not logged in")
            || result.localizedCaseInsensitiveContains("sign in")
            || result.localizedCaseInsensitiveContains("authenticate") {
            throw ClaudeCLIError.notAuthenticated
        }

        return parseUsageText(result)
    }

    static func parseUsageText(_ text: String) -> ClaudeCLIParsedUsage {
        var sessionPercent: Double?
        var weeklyPercent: Double?
        var weeklyResetText: String?
        var modelWeekly: [String: Double] = [:]

        let lines = text.components(separatedBy: .newlines)
        for line in lines {
            let trimmed = line.trimmingCharacters(in: .whitespaces)

            if let match = firstMatch(in: trimmed, pattern: #"Current session:\s*(\d+(?:\.\d+)?)%\s*used"#),
               let value = Double(match.captures[0]) {
                sessionPercent = value
                continue
            }

            if let match = firstMatch(in: trimmed, pattern: #"Current week \(all models\):\s*(\d+(?:\.\d+)?)%\s*used(?:\s*·\s*resets\s*(.+))?"#) {
                weeklyPercent = Double(match.captures[0])
                if match.captures.count > 1, !match.captures[1].isEmpty {
                    weeklyResetText = match.captures[1]
                }
                continue
            }

            if let match = firstMatch(in: trimmed, pattern: #"Current week \(([^)]+)\):\s*(\d+(?:\.\d+)?)%\s*used"#),
               let percent = Double(match.captures[1]) {
                modelWeekly[match.captures[0]] = percent
            }
        }

        return ClaudeCLIParsedUsage(
            sessionPercent: sessionPercent,
            weeklyPercent: weeklyPercent,
            weeklyResetText: weeklyResetText,
            modelWeekly: modelWeekly
        )
    }

    static func snapshot(from parsed: ClaudeCLIParsedUsage, fetchedAt: Date, isStale: Bool) -> ProviderUsageSnapshot {
        var values: [UsageMetric: MetricValue] = [:]

        if let sessionPercent = parsed.sessionPercent {
            values[.claudeFiveHourUtilization] = .percent(sessionPercent)
        }
        if let weeklyPercent = parsed.weeklyPercent {
            values[.claudeSevenDayUtilization] = .percent(weeklyPercent)
        }
        if let resetText = parsed.weeklyResetText {
            values[.claudeSevenDayResetsAt] = .text(resetText)
        }

        if let opus = parsed.modelWeekly["Opus"] {
            values[.claudeSevenDayOpusUtilization] = .percent(opus)
        }
        if let sonnet = parsed.modelWeekly["Sonnet"] {
            values[.claudeSevenDaySonnetUtilization] = .percent(sonnet)
        }
        for (model, percent) in parsed.modelWeekly where model != "Opus" && model != "Sonnet" {
            values[.claudeExtraUsage] = .text("\(model): \(Int(percent))%")
            break
        }

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

    private struct RegexMatch {
        let captures: [String]
    }

    private static func firstMatch(in text: String, pattern: String) -> RegexMatch? {
        guard let regex = try? NSRegularExpression(pattern: pattern, options: [.caseInsensitive]) else {
            return nil
        }
        let range = NSRange(text.startIndex..<text.endIndex, in: text)
        guard let match = regex.firstMatch(in: text, options: [], range: range) else {
            return nil
        }

        var captures: [String] = []
        for index in 1..<match.numberOfRanges {
            let captureRange = match.range(at: index)
            if captureRange.location != NSNotFound,
               let swiftRange = Range(captureRange, in: text) {
                captures.append(String(text[swiftRange]))
            } else {
                captures.append("")
            }
        }

        return RegexMatch(captures: captures)
    }
}