// SPDX-License-Identifier: AGPL-3.0-only
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0

import type { TrainingMethod } from "@/types/training";

/**
* VRAM estimation for model loading (4-bit quantization via bitsandbytes).
*
* Estimates the total driver-level VRAM (what nvidia-smi reports) to load a model in 4-bit with
* Unsloth / bitsandbytes, to check it fits the GPU before training.
*
* Formula: totalParams * 0.90 + 1.4 GB
*
* Calibrated against isolated Unsloth loads on RTX 5070 Ti (2026.2):
*   Qwen2.5-0.5B  (0.49B) : est 1.8 vs actual 1.86 GB  (-3%)
*   Llama-3.2-1B  (1.24B) : est 2.5 vs actual 2.54 GB   (-1%)
*   Llama-3.2-3B  (3.21B) : est 4.3 vs actual 4.40 GB   (-2%)
*   Llama-3.1-8B  (8.03B) : est 8.6 vs actual 8.14 GB   (+6%)
*/


/**
* Effective bytes per parameter for 4-bit weights at driver level. Raw bnb 4-bit is ~0.5, but
* embedding/lm_head stay fp16 and bnb adds per-block metadata, giving ~0.84-0.93; 0.9 is the mid. */
export const BNB_4BIT_LOADING_BYTES = 0.9;

/** Fixed overhead (GB) for the CUDA driver context and PyTorch runtime, independent of model
* size. Measured at 1.34-1.46 GB; we use 1.4. */
export const LOADING_OVERHEAD_GB = 1.4;

export type VramFitStatus = "fits" | "tight" | "exceeds";

/** Bytes per parameter at fp16/bf16 (LoRA, full FT). Theoretical (2 bytes); not yet calibrated,
* so real usage may run slightly higher (as 4-bit is 0.9 vs 0.5). */
export const FP16_LOADING_BYTES = 2.0;

function usesQuantizedLoading(
  method: TrainingMethod,
  modelId?: string,
): boolean {
  if (method === "qlora") return true;
  return method === "cpt" && (modelId ?? "").toLowerCase().includes("4bit");
}

/**
* Estimate VRAM (GB) to load a model with Unsloth. Bytes/param: QLoRA 4-bit bnb -> 0.90
* (calibrated); LoRA/Full/CPT fp16 -> 2.0. Formula: totalParams * bytesPerParam + 1.4 GB. */
export function estimateLoadingVram(
  totalParams: number,
  method: TrainingMethod = "qlora",
  modelId?: string,
): number {
  const bytesPerParam = usesQuantizedLoading(method, modelId)
    ? BNB_4BIT_LOADING_BYTES
    : FP16_LOADING_BYTES;
  const gb = (totalParams / 1e9) * bytesPerParam + LOADING_OVERHEAD_GB;
  return Math.round(gb * 10) / 10;
}

/** Check whether a model fits in available GPU VRAM: fits <= 75%; tight 75-100%; exceeds > 100%. */
export function checkVramFit(
  requiredGb: number,
  availableGb: number,
): VramFitStatus {
  if (availableGb <= 0) return requiredGb <= 0 ? "fits" : "exceeds";
  const ratio = requiredGb / availableGb;
  if (ratio <= 0.75) return "fits";
  if (ratio <= 1.0) return "tight";
  return "exceeds";
}

export interface ModelVramMapInput {
  id: string;
  totalParams?: number;
}

export interface ModelVramMapEntry {
  est: number;
  status: VramFitStatus | null;
}

export function buildModelVramMap(
  models: ModelVramMapInput[],
  method: TrainingMethod,
  gpu: { available: boolean; memoryTotalGb: number },
): Map<string, ModelVramMapEntry> {
  const map = new Map<string, ModelVramMapEntry>();
  for (const model of models) {
    if (!model.totalParams) {
      map.set(model.id, { est: 0, status: null });
      continue;
    }

    const est = estimateLoadingVram(model.totalParams, method, model.id);
    const status = gpu.available ? checkVramFit(est, gpu.memoryTotalGb) : null;
    map.set(model.id, { est, status });
  }
  return map;
}
