import * as os from 'os';

/**
 * System hardware specs surfaced to the rest of the app.
 *
 * All sizes are in **GiB** (1024^3 bytes). `safeModelBudgetGB` is a conservative
 * estimate of how much RAM is actually available for an LLM after the OS, the
 * editor, and the user's other apps take their cut — meant to be compared
 * against `IModelMemoryEstimator.estimate()` to warn the user before LM Studio
 * crashes on out-of-memory.
 */
export interface SystemSpecs {
    totalRamGB: number;
    freeRamGB: number;
    cpuModel: string;
    cpuCount: number;
    platform: NodeJS.Platform;
    arch: string;
    isAppleSilicon: boolean;
    safeModelBudgetGB: number;
    /** Human-readable one-line summary, useful for logs and toasts. */
    summary: string;
}

export interface ISystemSpecsProvider {
    /** Returns the current system specs. Implementations should cache. */
    get(): SystemSpecs;
}

export interface IModelMemoryEstimator {
    /**
     * Best-effort estimate of how many GB a given model identifier will need
     * to load. Returns a positive number; the caller decides what to do with it.
     */
    estimate(modelId: string): number;
}

/**
 * Production system-specs provider. Reads `os.totalmem()`, `os.cpus()`, etc.
 * and caches the result for the process lifetime — physical RAM does not
 * change at runtime.
 *
 * The Apple Silicon ratio (0.65) is more generous than other platforms (0.5)
 * because of unified memory: GPU inference shares the same pool, so a higher
 * fraction is realistically usable for the model.
 */
export class NodeSystemSpecsProvider implements ISystemSpecsProvider {
    private _cache: SystemSpecs | undefined;

    get(): SystemSpecs {
        if (this._cache) return this._cache;

        const totalRamGB = os.totalmem() / (1024 ** 3);
        const freeRamGB = os.freemem() / (1024 ** 3);
        const cpus = os.cpus() || [];
        const cpuModel = (cpus[0]?.model || 'unknown').replace(/\s+/g, ' ').trim();
        const platform = os.platform();
        const arch = os.arch();
        const isAppleSilicon =
            platform === 'darwin' && arch === 'arm64' && /Apple\s+M/i.test(cpuModel);

        const ratio = isAppleSilicon ? 0.65 : 0.5;
        const safeModelBudgetGB = Math.max(2, Math.floor(totalRamGB * ratio));

        const platformLabel = platform === 'darwin' ? 'macOS' : platform;
        const siliconLabel = isAppleSilicon ? ' (Apple Silicon)' : '';
        const summary =
            `${platformLabel} · ${arch}${siliconLabel} · RAM ${totalRamGB.toFixed(0)}GB ` +
            `· CPU ${cpuModel.slice(0, 40)} (${cpus.length} cores)`;

        this._cache = {
            totalRamGB,
            freeRamGB,
            cpuModel,
            cpuCount: cpus.length,
            platform,
            arch,
            isAppleSilicon,
            safeModelBudgetGB,
            summary,
        };
        return this._cache;
    }
}

/**
 * Heuristic estimator that derives memory cost from common patterns in model
 * identifiers (e.g. `gemma-2-9b-q4_K_M`).
 *
 * Approach:
 *   1. Pull the parameter count by matching `\d+B`. Default 7B if absent.
 *   2. Multiply by a per-param byte cost determined by quantization:
 *        - q4 / 4-bit: 0.6 GB/B (default)
 *        - q5 / 5-bit: 0.7
 *        - q6 / 6-bit: 0.8
 *        - q8 / 8-bit / fp8: 1.0
 *        - fp16 / bf16: 2.0
 *   3. Add a 1 GB overhead for KV cache + runtime.
 *
 * This is a rough estimate, but consistently within ±20% of LM Studio's actual
 * load footprint on the popular GGUF families — good enough to gate "your 16 GB
 * machine probably can't load a 70B fp16" before LM Studio crashes.
 */
export class HeuristicModelMemoryEstimator implements IModelMemoryEstimator {
    estimate(modelId: string): number {
        const id = (modelId || '').toLowerCase();
        const paramMatch = id.match(/(\d+(?:\.\d+)?)\s*b\b/);
        const paramCountB = paramMatch ? parseFloat(paramMatch[1]) : 7;

        let bytesPerParam = 0.6; // q4 default
        if (/q8|8bit|fp8/i.test(id)) bytesPerParam = 1.0;
        else if (/q6|6bit/i.test(id)) bytesPerParam = 0.8;
        else if (/q5|5bit/i.test(id)) bytesPerParam = 0.7;
        else if (/fp16|f16|bf16/i.test(id)) bytesPerParam = 2.0;

        return paramCountB * bytesPerParam + 1.0;
    }
}
