/** GPU profiler — detect hardware capabilities for model scheduling */

import { Effect } from "effect"
import * as Log from "@saeeol/core/util/log"
import { Process } from "@/util/process"
import type { GPUInfo, GPUProfile } from "./types"

const log = Log.create({ service: "local/gpu" })

/** Parse nvidia-smi output into structured GPU info */
function parseNvidiaSmi(stdout: string): GPUInfo[] {
  const gpus: GPUInfo[] = []
  // nvidia-smi --query-gpu=index,name,memory.total,memory.used,memory.free,compute_cap,driver_version,cuda_version --format=csv,noheader,nounits
  for (const line of stdout.trim().split("\n")) {
    const parts = line.split(",").map((s) => s.trim())
    if (parts.length < 5) continue
    gpus.push({
      index: Number(parts[0]),
      name: parts[1],
      vramTotalMB: Number(parts[2]),
      vramUsedMB: Number(parts[3]),
      vramFreeMB: Number(parts[4]),
      computeCapability: parts[5] || undefined,
      driverVersion: parts[6] || undefined,
      cudaVersion: parts[7] || undefined,
    })
  }
  return gpus
}

/** Profile all available GPUs via nvidia-smi */
export const profile: Effect.Effect<GPUProfile> = Effect.gen(function* () {
  const result = yield* Effect.tryPromise(() =>
    Process.run([
      "nvidia-smi",
      "--query-gpu=index,name,memory.total,memory.used,memory.free,compute_cap,driver_version,cuda_version",
      "--format=csv,noheader,nounits",
    ]),
  ).pipe(Effect.catch(() => Effect.succeed({ code: 1, stdout: Buffer.alloc(0), stderr: Buffer.alloc(0) } as Process.Result)))

  if (result.code !== 0) {
    log.info("no nvidia-smi found, GPU profiling unavailable")
    return {
      gpus: [],
      totalVRAMMB: 0,
      availableVRAMMB: 0,
      cudaAvailable: false,
    }
  }

  const gpus = parseNvidiaSmi(result.stdout.toString())
  return {
    gpus,
    totalVRAMMB: gpus.reduce((sum, g) => sum + g.vramTotalMB, 0),
    availableVRAMMB: gpus.reduce((sum, g) => sum + g.vramFreeMB, 0),
    cudaAvailable: gpus.length > 0,
  }
})

/** Estimate VRAM needed for a model (bytes) at a given quantization */
export function estimateVRAM(modelBytes: number, quantization: string): number {
  // Runtime overhead: ~20% for KV cache, context, etc.
  // GGUF files are already quantized — model file size ≈ VRAM needed + 20%
  const overhead = modelBytes * 0.2
  return Math.ceil((modelBytes + overhead) / (1024 * 1024)) // return MB
}

/** Find the best GPU for a model that needs `vramMB` megabytes */
export function findBestGPU(profile: GPUProfile, vramMB: number): GPUInfo | undefined {
  // Prefer GPU with most free VRAM that can fit the model
  const candidates = profile.gpus.filter((g) => g.vramFreeMB >= vramMB)
  if (candidates.length === 0) return undefined
  return candidates.sort((a, b) => b.vramFreeMB - a.vramFreeMB)[0]
}

/** Check if a model fits in available GPU memory */
export function canFit(profile: GPUProfile, vramMB: number): boolean {
  return profile.availableVRAMMB >= vramMB
}

/** Suggest quantization level based on available VRAM */
export function suggestQuantization(
  availableMB: number,
  modelSizeBytes: Record<string, number>,
): string | undefined {
  // Try from highest quality to lowest
  const order = ["fp16", "q8_0", "q6_k", "q5_k_m", "q5_k_s", "q4_k_m", "q4_k_s", "q3_k_m", "q3_k_s", "q2_k"]
  for (const q of order) {
    const size = modelSizeBytes[q]
    if (size !== undefined && estimateVRAM(size, q) <= availableMB) return q
  }
  return undefined
}
