#!/usr/bin/env node
// forge — tune and convert models for AMD Ryzen AI NPUs, one command.
//
//   forge tune <chats.jsonl> [--name grandma] [--base unsloth/Llama-3.2-1B-Instruct]
//              [--donor-tag llama3.2:1b] [--epochs 3]
//   forge convert <hf-repo> <file.gguf> [--name mymodel] [--donor-tag llama3.2:1b]
//   forge register            merge registry.json into FLM's list (UAC if needed)
//   forge list                stock models + forge customs with status
//   forge doctor              NPU stack / registry / port health checks
//   forge serve <tag>         flm serve passthrough (performance pmode)
'use strict';
const fs = require('fs');
const path = require('path');
const { spawnSync, spawn } = require('child_process');

const FORGE = __dirname;
const FLM_LIST = 'C:\\Program Files\\flm\\model_list.json';
const FLM_MODELS = path.join(process.env.USERPROFILE, 'Documents', 'flm', 'models');
const ENV = { ...process.env, PYTHONIOENCODING: 'utf-8' };

const argv = process.argv.slice(2);
const cmd = argv[0];
function flag(name, def) { const i = argv.indexOf('--' + name); return i > 0 ? argv[i + 1] : def; }
function run(exe, args, opts = {}) {
  const r = spawnSync(exe, args, { stdio: 'inherit', env: ENV, shell: true, ...opts });
  return r.status === 0;
}
function readJson(p) { try { return JSON.parse(fs.readFileSync(p, 'utf8').replace(/^/, '')); } catch (_) { return null; } }

function assemble(outName, tag, donorTag, donor) {
  return run('node', [path.join(FORGE, 'bin', 'assemble.js'), '--out', outName,
    '--donor', donor, '--tag', tag, '--donor-tag', donorTag]);
}

function donorFolderFor(donorTag) {
  const ml = readJson(FLM_LIST) || { models: {} };
  const [k, s] = (donorTag || 'llama3.2:1b').split(':');
  const e = ml.models[k] && ml.models[k][s];
  if (!e) { console.error('donor tag not found in FLM list: ' + donorTag); process.exit(1); }
  return e.name;
}

if (cmd === 'tune') {
  const data = argv[1];
  if (!data || !fs.existsSync(data)) { console.error('forge tune <chats.jsonl> — file not found'); process.exit(1); }
  const name = flag('name', path.basename(data, path.extname(data)));
  const base = flag('base', 'unsloth/Llama-3.2-1B-Instruct');
  const donorTag = flag('donor-tag', 'llama3.2:1b');
  const epochs = flag('epochs', '3');
  const outName = 'forge-' + name.toLowerCase().replace(/[^a-z0-9]+/g, '-');
  // stage data where the modal mount expects it
  const stage = path.join(FORGE, 'scratch', 'tune-data');
  fs.mkdirSync(stage, { recursive: true });
  fs.copyFileSync(data, path.join(stage, 'data.jsonl'));
  console.log(`forge tune: ${data} -> ${base} (${epochs} epochs) -> NPU as ${name}-forge`);
  if (!run('modal', ['run', path.join(FORGE, 'modal', 'tune_npu.py'),
    '--base', base, '--data-file', 'data.jsonl', '--out-name', outName, '--epochs', epochs])) {
    console.error('tune job failed — see output above'); process.exit(1);
  }
  if (!assemble(outName, name.toLowerCase() + '-forge:1b', donorTag, donorFolderFor(donorTag))) process.exit(1);
  console.log('\nDone. Run: forge register   then: flm run ' + name.toLowerCase() + '-forge:1b');
}

else if (cmd === 'convert') {
  const repo = argv[1], file = argv[2];
  if (!repo || !file) { console.error('forge convert <hf-repo> <file.gguf>'); process.exit(1); }
  const name = flag('name', file.replace(/\.gguf$/i, '').toLowerCase().replace(/[^a-z0-9]+/g, '-').slice(0, 30));
  const donorTag = flag('donor-tag', 'llama3.2:1b');
  const outName = 'forge-' + name;
  console.log(`forge convert: ${repo}/${file} -> NPU as ${name}-forge`);
  if (!run('modal', ['run', path.join(FORGE, 'modal', 'convert_q4nx.py'),
    '--repo-id', repo, '--filename', file, '--out-name', outName])) {
    console.error('convert job failed'); process.exit(1);
  }
  const size = donorTag.split(':')[1] || '1b';
  if (!assemble(outName, name + '-forge:' + size, donorTag, donorFolderFor(donorTag))) process.exit(1);
  console.log('\nDone. Run: forge register   then: flm run ' + name + '-forge:' + size);
}

else if (cmd === 'register') {
  const r = spawnSync('node', [path.join(FORGE, 'bin', 'register.js')], { encoding: 'utf8', env: ENV, shell: true });
  process.stdout.write(r.stdout || '');
  if (r.status !== 0) {
    console.log('(needs admin — launching UAC prompt via register-admin.bat)');
    spawn('cmd', ['/c', 'start', '', path.join(FORGE, 'register-admin.bat')], { detached: true, shell: false });
  }
}

else if (cmd === 'list') {
  run('flm', ['list', '--filter', 'installed']);
  const reg = readJson(path.join(FORGE, 'registry.json')) || { models: {} };
  const ml = readJson(FLM_LIST) || { models: {} };
  console.log('\nforge customs:');
  for (const [k, sizes] of Object.entries(reg.models)) {
    for (const [s, e] of Object.entries(sizes)) {
      const registered = !!(ml.models[k] && ml.models[k][s]);
      const folder = fs.existsSync(path.join(FLM_MODELS, e.name));
      console.log(`  ${k}:${s}  files:${folder ? 'OK' : 'MISSING'}  registered:${registered ? 'YES' : 'no — run forge register'}`);
    }
  }
}

else if (cmd === 'doctor') {
  console.log('— NPU device —');
  run('powershell', ['-NoProfile', '-Command', '"Get-PnpDevice -Class ComputeAccelerator | Select-Object Status, FriendlyName | Format-Table -HideTableHeaders"']);
  console.log('— NPU driver version (FLM needs >= 32.0.203.304, recommends .311) —');
  const dv = spawnSync('powershell', ['-NoProfile', '-Command', '"Get-PnpDevice -Class ComputeAccelerator | ForEach-Object { (Get-PnpDeviceProperty -InstanceId $_.InstanceId -KeyName DEVPKEY_Device_DriverVersion).Data }"'], { encoding: 'utf8', shell: true });
  const ver = (dv.stdout || '').trim().split('\n')[0] || 'unknown';
  const last = parseInt((ver.match(/(\d+)$/) || [])[1] || '0', 10);
  console.log(`  installed: ${ver}` + (last && last < 304 ? '  -> TOO OLD: update the AMD NPU driver (AMD Software / amd.com Ryzen AI driver), then re-run doctor' : last ? '  -> OK' : ''));
  console.log('— flm validate —');
  run('flm', ['validate']);
  console.log('— registry drift (customs lost to an flm update?) —');
  const reg = readJson(path.join(FORGE, 'registry.json')) || { models: {} };
  const ml = readJson(FLM_LIST) || { models: {} };
  let drift = 0;
  for (const [k, sizes] of Object.entries(reg.models)) for (const s of Object.keys(sizes)) {
    if (!(ml.models[k] && ml.models[k][s])) { console.log(`  LOST: ${k}:${s} — run forge register`); drift++; }
  }
  if (!drift) console.log('  all forge customs registered');
  console.log('— port 52625 —');
  const ns = spawnSync('netstat', ['-ano'], { encoding: 'utf8' });
  console.log((ns.stdout || '').split('\n').some(l => l.includes(':52625') && l.includes('LISTENING')) ? '  flm serve is UP' : '  free (flm not serving)');
}

else if (cmd === 'serve') {
  const tag = argv[1];
  if (!tag) { console.error('forge serve <tag>'); process.exit(1); }
  run('flm', ['serve', tag, '--pmode', 'performance']);
}

else {
  console.log(fs.readFileSync(__filename, 'utf8').split('\n').slice(1, 11).map(l => l.replace(/^\/\/ ?/, '')).join('\n'));
}