#!/usr/bin/env node
// assemble.js — local leg of forge convert: pull the converted .q4nx from the
// Modal volume, build the FLM model folder (donor config/tokenizer from the
// same-architecture stock model), and add the registry.json entry.
//
//   node bin/assemble.js --out forge-llama32-1b --donor Llama-3.2-1B-NPU2 \
//        --tag llama32-forge:1b --donor-tag llama3.2:1b
'use strict';
const fs = require('fs');
const path = require('path');
const { spawnSync } = require('child_process');

const FORGE = path.join(__dirname, '..');
const FLM_MODELS = path.join(process.env.USERPROFILE, 'Documents', 'flm', 'models');
const FLM_LIST = 'C:\\Program Files\\flm\\model_list.json';

function arg(name, def) {
  const i = process.argv.indexOf('--' + name);
  return i > 0 ? process.argv[i + 1] : def;
}

function main() {
  const outName = arg('out');                       // volume folder name
  const donor = arg('donor');                       // stock model folder to copy config/tokenizer from
  const tag = arg('tag');                           // e.g. llama32-forge:1b
  const donorTag = arg('donor-tag');                // e.g. llama3.2:1b (registry donor entry)
  if (!outName || !donor || !tag || !donorTag) {
    console.error('usage: node bin/assemble.js --out <volumeName> --donor <StockFolder> --tag <key:size> --donor-tag <stockKey:size>');
    process.exit(1);
  }
  const forgeName = 'Forge-' + outName.replace(/^forge-/, '') + '-NPU2';
  const dest = path.join(FLM_MODELS, forgeName);
  fs.mkdirSync(dest, { recursive: true });

  // 1. download q4nx + tokenizer.json (the converter extracts it from GGUF)
  console.log('[1/3] downloading from Modal volume npu-forge-out/' + outName);
  for (const f of ['model.q4nx', 'tokenizer.json']) {
    const r = spawnSync('modal', ['volume', 'get', 'npu-forge-out', outName + '/' + f, path.join(dest, f), '--force'],
      { encoding: 'utf8', env: { ...process.env, PYTHONIOENCODING: 'utf-8' }, shell: true, timeout: 900000 });
    if (!fs.existsSync(path.join(dest, f))) {
      console.error('download failed for ' + f + ':\n' + (r.stderr || r.stdout || '').slice(-800));
      process.exit(1);
    }
    console.log('  ' + f + ': ' + (fs.statSync(path.join(dest, f)).size / 1e6).toFixed(1) + ' MB');
  }

  // 2. donor config files (same architecture => identical config works)
  console.log('[2/3] copying config.json + tokenizer_config.json from donor ' + donor);
  const donorDir = path.join(FLM_MODELS, donor);
  for (const f of ['config.json', 'tokenizer_config.json']) {
    fs.copyFileSync(path.join(donorDir, f), path.join(dest, f));
  }

  // 3. registry entry from the donor's model_list entry
  console.log('[3/3] adding registry.json entry ' + tag);
  const ml = JSON.parse(fs.readFileSync(FLM_LIST, 'utf8').replace(/^/, ''));
  const [dKey, dSize] = donorTag.split(':');
  const donorEntry = ml.models[dKey] && ml.models[dKey][dSize];
  if (!donorEntry) { console.error('donor tag not in model_list: ' + donorTag); process.exit(1); }
  const entry = JSON.parse(JSON.stringify(donorEntry));
  entry.name = forgeName; entry.url = ''; entry.file_url = '';
  entry.modified_at = new Date().toISOString();
  const regPath = path.join(FORGE, 'registry.json');
  const reg = fs.existsSync(regPath) ? JSON.parse(fs.readFileSync(regPath, 'utf8').replace(/^/, '')) : { models: {} };
  const [key, size] = tag.split(':');
  reg.models[key] = reg.models[key] || {};
  reg.models[key][size] = entry;
  fs.writeFileSync(regPath, JSON.stringify(reg, null, 2));

  console.log('\nAssembled: ' + dest);
  console.log('Registered in forge registry as ' + tag + ' — now double-click register-admin.bat, then:');
  console.log('  flm run ' + tag);
}

main();