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

import Foundation
import MLX

/// RAM feasibility verdicts for inference and training on *this* Mac.
///
/// Formulas per docs/MODELS.md §3 and docs/TRAINING-RESEARCH.md §6:
/// - weights: bits-per-param + group-scale overhead (4-bit gs64 affine
///   = 4.5 bits ≈ 0.5625 B/param, confirmed against real repo sizes)
/// - inference ≈ weights + KV cache + ~20% overhead
/// - LoRA (bf16 base) ≈ 2 B/param + activations; QLoRA ≈ quantized weights
///   + activations; full FT ≈ 8 B/param (Adam) + activations
/// - the GPU working-set ceiling comes from the device, not a constant:
///   `GPU.deviceInfo().maxRecommendedWorkingSetSize`
enum MemoryVerdict: String, Codable, Sendable {
    case comfortable
    case tight
    case wontFit

    var displayName: String {
        switch self {
        case .comfortable: "Fits"
        case .tight: "Tight"
        case .wontFit: "Won't fit"
        }
    }
}

enum MemoryAdvisor {

    /// Physical unified memory of this Mac.
    static var physicalMemory: Int64 { Int64(ProcessInfo.processInfo.physicalMemory) }

    /// The GPU-visible working-set ceiling reported by Metal via MLX.
    static var recommendedWorkingSet: Int64 {
        let info = GPU.deviceInfo()
        let value = info.maxRecommendedWorkingSetSize
        return value > 0 ? Int64(value) : physicalMemory * 3 / 4
    }

    /// Effective bytes per parameter for a given quantization
    /// (group-scale overhead included: bits + 32/groupSize extra bits).
    static func bytesPerParameter(quantization: QuantizationInfo?) -> Double {
        guard let q = quantization else { return 2.0 } // bf16
        let effectiveBits = Double(q.bits) + 32.0 / Double(q.groupSize)
        return effectiveBits / 8.0
    }

    /// Estimated KV-cache bytes for a typical mid-size GQA model at the given
    /// context. Uses conservative defaults when architecture details are
    /// unknown (calibrated in Phase 7).
    static func estimatedKVCache(parameterCount: Int64?, contextLength: Int) -> Int64 {
        // Scale a measured anchor: ~1.2 GB at 8k ctx for an 8B GQA model
        // (docs/MLX-RESEARCH.md §7), linear in context and sublinear in params.
        let params = Double(parameterCount ?? 8_000_000_000)
        let anchor = 1.2 * 1_073_741_824.0
        let scale = (params / 8_000_000_000.0).squareRoot()
        return Int64(anchor * scale * Double(contextLength) / 8192.0)
    }

    // MARK: - Inference

    static func inferenceBytes(for model: LocalModel, contextLength: Int = 8192) -> Int64 {
        let weights = Double(model.weightsSize)
        let kv = Double(estimatedKVCache(parameterCount: model.parameterCount, contextLength: contextLength))
        return Int64(weights * 1.2 + kv)
    }

    static func inferenceVerdict(for model: LocalModel, contextLength: Int = 8192) -> MemoryVerdict {
        verdict(needed: inferenceBytes(for: model, contextLength: contextLength))
    }

    // MARK: - Training

    static func trainingBytes(for model: LocalModel, method: FineTuneMethod, params: HyperParams) -> Int64 {
        let count = Double(model.parameterCount ?? 0)
        guard count > 0 else { return .max }

        // Activation estimate: batch × seq drives it; grad-checkpoint ~halves it.
        var activations = Double(params.batchSize) * Double(params.maxSeqLength) / (4.0 * 2048.0)
            * 2.0 * 1_073_741_824.0
        if params.gradCheckpoint { activations *= 0.55 }

        let base: Double
        switch method {
        case .lora, .dora:
            base = count * 2.0 // bf16 base held in memory
        case .qlora:
            base = Double(model.weightsSize) // quantized base stays quantized
        case .full:
            base = count * 8.0 // weights + grads + 2 Adam moments (bf16)
        }
        // Adapter weights + their optimizer state are negligible (MB-scale).
        return Int64(base + activations)
    }

    static func trainingVerdict(for model: LocalModel, method: FineTuneMethod, params: HyperParams) -> MemoryVerdict {
        verdict(needed: trainingBytes(for: model, method: method, params: params))
    }

    /// Suggested remedies when a config won't fit (charter 3.B).
    static func suggestions(for model: LocalModel, method: FineTuneMethod, params: HyperParams) -> [String] {
        var out: [String] = []
        if method == .lora || method == .dora {
            out.append("Use QLoRA (train on a 4-bit base) — largest single saving")
        }
        if params.batchSize > 1 { out.append("Reduce batch size (\(params.batchSize) → \(max(1, params.batchSize / 2))) and use gradient accumulation") }
        if params.numLayers == -1 || params.numLayers > 8 { out.append("Adapt fewer layers (--num-layers 8 or 4)") }
        if !params.gradCheckpoint { out.append("Enable gradient checkpointing") }
        if params.maxSeqLength > 1024 { out.append("Lower max sequence length or pre-split long samples") }
        out.append("Choose a smaller base model")
        return out
    }

    // MARK: - Verdict core

    private static func verdict(needed: Int64) -> MemoryVerdict {
        let ceiling = recommendedWorkingSet
        if needed <= Int64(Double(ceiling) * 0.75) { return .comfortable }
        if needed <= ceiling { return .tight }
        return .wontFit
    }
}