import { runProcess } from '../process.js';

const NVIDIA_QUERY = [
  '--query-gpu=index,name,memory.used,memory.total,utilization.gpu',
  '--format=csv,noheader,nounits'
];

function parseGpuLine(line) {
  const [index, name, memoryUsed, memoryTotal, utilization] = line.split(',').map((part) => part.trim());
  return {
    index: Number(index),
    name,
    memoryUsedMiB: Number(memoryUsed),
    memoryTotalMiB: Number(memoryTotal),
    gpuUtilizationPercent: Number(utilization)
  };
}

export function parseNvidiaSmiCsv(output) {
  return output
    .split('\n')
    .map((line) => line.trim())
    .filter(Boolean)
    .map(parseGpuLine)
    .filter((gpu) => Number.isFinite(gpu.index));
}

export async function collectGpuStatus(config) {
  const hosts = config.gpuHosts || [];
  const checks = hosts.map(async (host) => {
    const result = await runProcess('ssh', [
      '-o', 'BatchMode=yes',
      '-o', `ConnectTimeout=${Math.max(1, Math.ceil((config.sshTimeoutMs || 10000) / 1000))}`,
      host,
      'nvidia-smi',
      ...NVIDIA_QUERY
    ], { timeoutMs: config.sshTimeoutMs || 10000 });

    return {
      host,
      ok: result.ok,
      checkedAt: new Date().toISOString(),
      gpus: result.ok ? parseNvidiaSmiCsv(result.stdout) : [],
      error: result.ok ? null : result.stderr || result.stdout || 'Unable to query GPU status.'
    };
  });

  return Promise.all(checks);
}