//
//  MemoryAdvisor.swift
//  Zyquo Local
//
//  Author: Simon-Pierre Boucher
//  Mail: contact@spboucher.ai
//

import Foundation
import MLX

/// RAM estimation and fits/tight/won't-fit verdicts for THIS Mac.
/// Thresholds follow docs/MLX-RESEARCH.md §7: the macOS GPU-wired ceiling is
/// ~70–75 % of unified memory, so weights ≤60 % is comfortable, ≤75 % tight.
enum MemoryAdvisor {
    enum Verdict: String, Codable, Sendable {
        case fits
        case tight
        case tooLarge

        var label: String {
            switch self {
            case .fits: "Fits"
            case .tight: "Tight"
            case .tooLarge: "Too large"
            }
        }
    }

    /// Physical unified memory of this Mac in bytes.
    static var physicalMemoryBytes: UInt64 { HardwareGate.physicalMemory }

    /// Estimated RAM needed to run a model: weights + KV cache + overhead.
    /// weights×1.2 covers activations/prefill working set; +1.5 GB covers a
    /// useful KV cache at moderate context.
    static func estimatedRuntimeBytes(weightsBytes: Int64) -> Int64 {
        Int64(Double(weightsBytes) * 1.2) + 1_500_000_000
    }

    /// Verdict for a model of the given on-disk weights size.
    static func verdict(weightsBytes: Int64) -> Verdict {
        let ram = Double(physicalMemoryBytes)
        let weights = Double(weightsBytes)
        if weights <= ram * 0.60 { return .fits }
        if weights <= ram * 0.75 { return .tight }
        return .tooLarge
    }

    /// Verdict for running two models side by side (Compare mode gate).
    static func verdict(combinedWeightsBytes: Int64) -> Verdict {
        verdict(weightsBytes: combinedWeightsBytes)
    }

    /// Live MLX memory snapshot (bytes actively used by the loaded model).
    static var activeMemoryBytes: Int { MLX.Memory.activeMemory }
    static var peakMemoryBytes: Int { MLX.Memory.peakMemory }

    static func resetPeakMemory() { MLX.GPU.resetPeakMemory() }

    /// Returns cached MLX buffers to the OS — called after unloading a model.
    static func reclaimMemory() { MLX.Memory.clearCache() }
}