#!/usr/bin/env node

const { execSync } = require('child_process');
const fs = require('fs');
const path = require('path');

const DEFAULT_CONFIG_PATH = '/etc/amd-gpu-temp-control.json';
const DEFAULT_MIN_TEMP = 70;
const DEFAULT_MAX_TEMP = 80;
const DEFAULT_MIN_TDP_PERCENT = 10;
const DEFAULT_MAX_TDP_PERCENT = 100;
const DEFAULT_TDP_STEP_PERCENT = 10;
const DEFAULT_CHECK_INTERVAL_MS = 2000;

function loadConfig(configPath) {
  if (!fs.existsSync(configPath)) {
    return null;
  }
  const raw = fs.readFileSync(configPath, 'utf-8');
  return JSON.parse(raw);
}

function mergeConfig(config) {
  return {
    minTemp: config.minTemp ?? DEFAULT_MIN_TEMP,
    maxTemp: config.maxTemp ?? DEFAULT_MAX_TEMP,
    minTdpPercent: config.minTdpPercent ?? DEFAULT_MIN_TDP_PERCENT,
    maxTdpPercent: config.maxTdpPercent ?? DEFAULT_MAX_TDP_PERCENT,
    tdpStepPercent: config.tdpStepPercent ?? DEFAULT_TDP_STEP_PERCENT,
    checkIntervalMs: config.checkIntervalMs ?? DEFAULT_CHECK_INTERVAL_MS,
  };
}

function getCardIndices() {
  const entries = fs.readdirSync('/sys/class/drm');
  const indices = [];
  for (const name of entries) {
    const match = name.match(/^card(\d+)$/);
    if (match) {
      indices.push(parseInt(match[1], 10));
    }
  }
  return indices.length > 0 ? indices : [0];
}

function findHwmonDirs(cardNumber) {
  const hwmonDirs = [];
  const baseDir = `/sys/class/drm/card${cardNumber}/device/hwmon`;

  if (fs.existsSync(baseDir)) {
    for (const entry of fs.readdirSync(baseDir)) {
      if (entry.startsWith('hwmon')) {
        hwmonDirs.push(path.join(baseDir, entry));
      }
    }
  }

  const altBase = `/sys/class/drm/card${cardNumber}/device`;
  if (fs.existsSync(altBase)) {
    for (const entry of fs.readdirSync(altBase)) {
      if (entry.startsWith('hwmon')) {
        const dirPath = path.join(altBase, entry);
        if (!hwmonDirs.includes(dirPath)) {
          hwmonDirs.push(dirPath);
        }
      }
    }
  }

  return hwmonDirs;
}

function getTemperature(cardNumber) {
  const hwmonDirs = findHwmonDirs(cardNumber);
  if (hwmonDirs.length === 0) return null;

  let maxTemp = null;

  for (const hwmonDir of hwmonDirs) {
    try {
      const entries = fs.readdirSync(hwmonDir);
      for (const entry of entries) {
        const match = entry.match(/^temp\d+_input$/);
        if (match) {
          const raw = fs.readFileSync(path.join(hwmonDir, entry), 'utf-8').trim();
          const value = parseInt(raw, 10);
          if (!isNaN(value)) {
            const tempC = value / 1000;
            if (maxTemp === null || tempC > maxTemp) {
              maxTemp = tempC;
            }
          }
        }
      }
    } catch {
    }
  }

  return maxTemp;
}

function getDeviceName(cardNumber) {
  try {
    const namePath = `/sys/class/drm/card${cardNumber}/device/name`;
    if (fs.existsSync(namePath)) {
      return fs.readFileSync(namePath, 'utf-8').trim();
    }
  } catch { /* ignore */ }

  try {
    const slotPath = `/sys/class/drm/card${cardNumber}/device/uevent`;
    if (fs.existsSync(slotPath)) {
      const lines = fs.readFileSync(slotPath, 'utf-8').trim().split('\n');
      for (const line of lines) {
        const match = line.match(/^PCI_SLOT_NAME=(.*)$/);
        if (match) {
          const slot = match[1].trim();
          try {
            const out = execSync(`lspci -s ${slot} 2>/dev/null`, { encoding: 'utf-8' }).trim();
            if (out) return out.split(/\s{3,}/).pop().trim();
          } catch { /* ignore */ }
        }
      }
    }
  } catch { /* ignore */ }

  return `card${cardNumber}`;
}

function getTdpLimit(cardNumber) {
  try {
    const out = execSync(`amd-smi static --json -g ${cardNumber}`, { encoding: 'utf-8', stdio: 'pipe' }).trim();
    const data = JSON.parse(out);
    const gpuData = data.gpu_data[0];
    if (gpuData && gpuData.limit && gpuData.limit.ppt0 && gpuData.limit.ppt0.max_power_limit) {
      return gpuData.limit.ppt0.max_power_limit.value;
    }
  } catch { /* ignore */ }
  return 100;
}

function setTdpForCard(cardNumber, power) {
  execSync(`amd-smi set -g ${cardNumber} -o ppt0 ${power}`, { encoding: 'utf-8', stdio: 'pipe' });
}

function parseArgs(argv) {
  let configPath = null;

  for (let i = 2; i < argv.length; i++) {
    switch (argv[i]) {
      case '--config':
      case '-c':
        configPath = argv[++i];
        break;
      case '--help':
      case '-h':
        console.log('Usage: node amd-gpu-temp-control.js [options]');
        console.log('');
        console.log('Options:');
        console.log('  --config, -c <path>   Config file path (default: /etc/amd-gpu-temp-control.json)');
        console.log('  --help, -h            Show this help message');
        process.exit(0);
        break;
    }
  }

  return configPath || DEFAULT_CONFIG_PATH;
}

function calculateTargetTdpPercent(temp, minTemp, maxTemp, minTdpPercent, maxTdpPercent) {
  if (temp <= minTemp) return maxTdpPercent;
  if (temp >= maxTemp) return minTdpPercent;

  const ratio = (temp - minTemp) / (maxTemp - minTemp);
  return Math.round(maxTdpPercent - ratio * (maxTdpPercent - minTdpPercent));
}

function wattsFromPercent(percent, maxTdp) {
  return Math.round((percent / 100) * maxTdp);
}

function log(message) {
  const timestamp = new Date().toISOString();
  console.log(`[${timestamp}] ${message}`);
}

function main() {
  const configPath = parseArgs(process.argv);
  const rawConfig = loadConfig(configPath);
  const config = mergeConfig(rawConfig || {});

  const cardIndices = getCardIndices();

  log('AMD GPU Temperature Control');
  log(`Temperature range: ${config.minTemp}°C - ${config.maxTemp}°C`);
  log(`TDP range: ${config.minTdpPercent}% - ${config.maxTdpPercent}%`);
  log(`TDP step: ${config.tdpStepPercent}%`);
  log(`Check interval: ${config.checkIntervalMs}ms`);
  log(`Config: ${configPath}`);
  log(`Monitoring ${cardIndices.length} GPU(s)`);
  log('Press Ctrl+C to stop');
  log('');

  const gpuStates = cardIndices.map(cardNumber => ({
    cardNumber,
    deviceName: getDeviceName(cardNumber),
    maxTdp: getTdpLimit(cardNumber),
    currentTdpPercent: null,
  }));

  for (const state of gpuStates) {
    log(`GPU ${state.cardNumber} [${state.deviceName}]: max TDP ${state.maxTdp}W`);
  }

  if (gpuStates.length === 0) {
    log('No GPUs found');
    process.exit(1);
  }

  let iteration = 0;

  const interval = setInterval(() => {
    iteration++;

    for (const gpuState of gpuStates) {

      try {
        const temp = getTemperature(gpuState.cardNumber);
        if (temp === null) {
          log(`GPU ${gpuState.cardNumber} [${gpuState.deviceName}]: Cannot read temperature`);
          continue;
        }

        const inRange = temp >= config.minTemp && temp <= config.maxTemp;

        if (!inRange) {
          const targetTdpPercent = calculateTargetTdpPercent(
            temp,
            config.minTemp,
            config.maxTemp,
            config.minTdpPercent,
            config.maxTdpPercent
          );

          let shouldApply = false;

          if (gpuState.currentTdpPercent === null) {
            gpuState.currentTdpPercent = targetTdpPercent;
            shouldApply = true;
          } else if (targetTdpPercent !== gpuState.currentTdpPercent) {
            const diff = targetTdpPercent - gpuState.currentTdpPercent;

            if (Math.abs(diff) >= config.tdpStepPercent) {
              gpuState.currentTdpPercent += Math.sign(diff) * config.tdpStepPercent;
            } else {
              gpuState.currentTdpPercent = targetTdpPercent;
            }
            shouldApply = true;
          }

          if (shouldApply) {
            const power = wattsFromPercent(gpuState.currentTdpPercent, gpuState.maxTdp);
            setTdpForCard(gpuState.cardNumber, power);
            log(`GPU ${gpuState.cardNumber} [${gpuState.deviceName}]: ${temp.toFixed(1)}°C | TDP: ${gpuState.currentTdpPercent}% (${power}W)`);
          }
        } else if (gpuState.currentTdpPercent !== null && iteration % 10 === 0) {
          const currentWatts = wattsFromPercent(gpuState.currentTdpPercent, gpuState.maxTdp);
          log(`GPU ${gpuState.cardNumber} [${gpuState.deviceName}]: ${temp.toFixed(1)}°C | TDP: ${gpuState.currentTdpPercent}% (${currentWatts}W) (stable)`);
        }
      } catch (err) {
        log(`GPU ${gpuState.cardNumber} [${gpuState.deviceName}]: Error reading temperature - ${err.message}`);
      }
    }
  }, config.checkIntervalMs);

  process.on('SIGINT', () => {
    clearInterval(interval);
    log('');
    log('Stopped.');
    process.exit(0);
  });
}

main();