markbaseengine/Sources/MarkBaseServer/PerformanceMetrics.swift

import Foundation

// ─────────────────────────────────────────────────────────────
// Performance Metrics API Models
// Prometheus Compatible Format
// ─────────────────────────────────────────────────────────────

/// Performance metrics collector
public final class MetricsCollector {
    public nonisolated(unsafe) static let shared = MetricsCollector()
    
    private var counters: [String: CounterMetric] = [:]
    private var histograms: [String: MutableHistogram] = [:]
    private let lock = NSLock()
    
    private init() {}
    
    /// Record a request
    public func recordRequest(duration: TimeInterval, tokens: Int, model: String) {
        lock.lock()
        defer { lock.unlock() }
        
        // Request count
        incrementCounter(name: "markbase_requests_total", labels: ["model": model])
        
        // Request duration
        observeHistogram(name: "markbase_request_duration_seconds", value: duration, labels: ["model": model])
        
        // Tokens processed
        incrementCounter(name: "markbase_tokens_total", labels: ["model": model, "type": "output"], value: Double(tokens))
    }
    
    /// Record a token generation
    public func recordToken(tokens: Int, duration: TimeInterval, model: String) {
        lock.lock()
        defer { lock.unlock() }
        
        let tokensPerSecond = Double(tokens) / duration
        observeHistogram(name: "markbase_tokens_per_second", value: tokensPerSecond, labels: ["model": model])
    }
    
    /// Record an error
    public func recordError(type: String, model: String) {
        lock.lock()
        defer { lock.unlock() }
        
        incrementCounter(name: "markbase_errors_total", labels: ["model": model, "type": type])
    }
    
    /// Get Prometheus format metrics
    public func getPrometheusMetrics() -> String {
        lock.lock()
        defer { lock.unlock() }
        
        var output = ""
        
        for (name, counter) in counters {
            output += counter.toPrometheus(name: name)
        }
        
        for (name, histogram) in histograms {
            output += histogram.toPrometheus(name: name)
        }
        
        return output
    }
    
    // ─────────────────────────────────────────────────────────────
    // Private helpers
    // ─────────────────────────────────────────────────────────────
    
    private func incrementCounter(name: String, labels: [String: String], value: Double = 1.0) {
        if counters[name] == nil {
            counters[name] = CounterMetric()
        }
        counters[name]?.increment(by: value, labels: labels)
    }
    
    private func observeHistogram(name: String, value: Double, labels: [String: String]) {
        if histograms[name] == nil {
            histograms[name] = MutableHistogram()
        }
        histograms[name]?.observe(value, labels: labels)
    }
}

// ─────────────────────────────────────────────────────────────
// Metric types
// ─────────────────────────────────────────────────────────────

protocol Metric {
    func toPrometheus(name: String) -> String
}

struct CounterMetric: Metric {
    var values: [[String: String]: Double] = [:]
    
    mutating func increment(by value: Double, labels: [String: String]) {
        values[labels, default: 0] += value
    }
    
    func toPrometheus(name: String) -> String {
        var output = "# TYPE \(name) counter\n"
        
        for (labels, value) in values {
            let labelsStr = labels.map { "\($0.key)=\"\($0.value)\"" }.joined(separator: ",")
            output += "\(name){\(labelsStr)} \(value)\n"
        }
        
        return output + "\n"
    }
}

struct HistogramMetric: Metric {
    var counts: [[String: String]: Int] = [:]
    var sums: [[String: String]: Double] = [:]
    
    mutating func observe(_ value: Double, labels: [String: String]) {
        counts[labels, default: 0] += 1
        sums[labels, default: 0] += value
    }
    
    func toPrometheus(name: String) -> String {
        var output = "# TYPE \(name) histogram\n"
        
        for (labels, count) in counts {
            let sum = sums[labels] ?? 0
            let labelsStr = labels.map { "\($0.key)=\"\($0.value)\"" }.joined(separator: ",")
            output += "\(name)_count{\(labelsStr)} \(count)\n"
            output += "\(name)_sum{\(labelsStr)} \(sum)\n"
        }
        
        return output + "\n"
    }
}

// Wrapper for mutable histogram
final class MutableHistogram {
    var histogram: HistogramMetric
    
    init() {
        self.histogram = HistogramMetric()
    }
    
    func observe(_ value: Double, labels: [String: String]) {
        histogram.observe(value, labels: labels)
    }
    
    func toPrometheus(name: String) -> String {
        histogram.toPrometheus(name: name)
    }
}