TokenLens/NativeTokenLens/Sources/TokenLensCore/TokenUsageParser.swift

import Foundation

public enum TokenUsageParser {
    public static func parseClaudeLine(_ line: String, sourcePath: String) -> UsageRecord? {
        guard let object = parseJSONObject(line) else { return nil }
        let message = object["message"] as? [String: Any]
        let usage = message?["usage"] as? [String: Any] ?? object["usage"] as? [String: Any]
        guard let usage else { return nil }
        guard let timestamp = parseDate(object["timestamp"] ?? object["createdAt"] ?? object["date"]) else { return nil }

        return UsageRecord(
            provider: .claude,
            source: sourceLabel(sourcePath: sourcePath, provider: .claude),
            sourcePath: sourcePath,
            timestamp: timestamp,
            sessionID: stringValue(object["sessionId"]),
            model: stringValue(message?["model"]) == "" ? "Claude" : stringValue(message?["model"]),
            inputTokens: intValue(usage["input_tokens"]),
            outputTokens: intValue(usage["output_tokens"]),
            cachedTokens: intValue(usage["cache_read_input_tokens"]) + intValue(usage["cache_creation_input_tokens"]),
            reasoningTokens: 0,
            toolTokens: 0
        )
    }

    public static func parseCodexLine(_ line: String, sourcePath: String) -> UsageRecord? {
        guard let object = parseJSONObject(line),
              let payload = object["payload"] as? [String: Any],
              let info = payload["info"] as? [String: Any],
              let usage = info["last_token_usage"] as? [String: Any],
              let timestamp = parseDate(object["timestamp"])
        else { return nil }

        return UsageRecord(
            provider: .codex,
            source: sourceLabel(sourcePath: sourcePath, provider: .codex),
            sourcePath: sourcePath,
            timestamp: timestamp,
            sessionID: stringValue(payload["id"] ?? payload["thread_id"]),
            model: stringValue(payload["model"]) == "" ? "Codex" : stringValue(payload["model"]),
            inputTokens: intValue(usage["input_tokens"]),
            outputTokens: intValue(usage["output_tokens"]),
            cachedTokens: intValue(usage["cached_input_tokens"]),
            reasoningTokens: intValue(usage["reasoning_output_tokens"]),
            toolTokens: 0
        )
    }

    public static func parseGeminiLine(_ line: String, sourcePath: String) -> UsageRecord? {
        guard let object = parseJSONObject(line),
              let tokens = object["tokens"] as? [String: Any],
              let timestamp = parseDate(object["timestamp"])
        else { return nil }

        return UsageRecord(
            provider: .gemini,
            source: sourceLabel(sourcePath: sourcePath, provider: .gemini),
            sourcePath: sourcePath,
            timestamp: timestamp,
            sessionID: stringValue(object["id"]),
            model: stringValue(object["model"]) == "" ? "Gemini" : stringValue(object["model"]),
            inputTokens: intValue(tokens["input"]),
            outputTokens: intValue(tokens["output"]),
            cachedTokens: intValue(tokens["cached"]),
            reasoningTokens: intValue(tokens["thoughts"]),
            toolTokens: intValue(tokens["tool"]),
            explicitTotal: intValue(tokens["total"]) == 0 ? nil : intValue(tokens["total"])
        )
    }

    static func parseDate(_ value: Any?) -> Date? {
        guard let text = value as? String else { return nil }
        let isoFormatter = ISO8601DateFormatter()
        isoFormatter.formatOptions = [.withInternetDateTime, .withFractionalSeconds]
        let fallbackISOFormatter = ISO8601DateFormatter()
        fallbackISOFormatter.formatOptions = [.withInternetDateTime]
        return isoFormatter.date(from: text) ?? fallbackISOFormatter.date(from: text)
    }

    private static func parseJSONObject(_ line: String) -> [String: Any]? {
        guard let data = line.data(using: .utf8) else { return nil }
        return (try? JSONSerialization.jsonObject(with: data)) as? [String: Any]
    }

    private static func stringValue(_ value: Any?) -> String {
        value as? String ?? ""
    }

    private static func intValue(_ value: Any?) -> Int {
        if let int = value as? Int { return int }
        if let double = value as? Double { return Int(double) }
        if let string = value as? String { return Int(string) ?? 0 }
        return 0
    }

    private static func sourceLabel(sourcePath: String, provider: Provider) -> String {
        switch provider {
        case .claude: return "projects/*.jsonl"
        case .codex: return "sessions/*.jsonl"
        case .gemini: return "tmp/chats/*.jsonl"
        case .hermes: return "state.db/sessions"
        case .opencode: return "opencode.db/session"
        }
    }
}