#!/usr/bin/env node
/**
 * generate-weapon-images.js
 *
 * Generates AI artwork for each Deathwatch weapon/armour card and saves
 * them to public/weapon-images/.
 *
 * Supports two providers:
 *   --provider openai   (DALL-E 3, ~$0.04/image)
 *   --provider gemini   (Imagen 3, free tier available)
 *
 * Usage:
 *   GOOGLE_API_KEY=... node scripts/generate-weapon-images.js --provider gemini
 *   GOOGLE_API_KEY=... node scripts/generate-weapon-images.js --provider gemini --limit 3
 *   OPENAI_API_KEY=... node scripts/generate-weapon-images.js --provider openai --limit 3
 *   Add --force to regenerate already-existing images
 */

const fetch = require('node-fetch');
const fs = require('fs');
const path = require('path');

// ── CLI args ──────────────────────────────────────────────────────────────
const args = process.argv.slice(2);

const PROVIDER = (() => {
  const i = args.indexOf('--provider');
  return i !== -1 ? args[i + 1] : 'huggingface';
})();
const LIMIT = (() => {
  const i = args.indexOf('--limit');
  return i !== -1 ? parseInt(args[i + 1], 10) : Infinity;
})();
const FORCE = args.includes('--force');

// ── Config ────────────────────────────────────────────────────────────────
const OUTPUT_DIR  = path.join(__dirname, '../public/weapon-images');
const MANIFEST    = path.join(OUTPUT_DIR, 'manifest.json');
const WEAPONS_API = 'http://localhost:5000/api/weapons';

// Rate limits
const DELAY = {
  openai: 13_000,  // DALL-E 3: 5 req/min on tier-1
  gemini: 4_000,   // Imagen 3: 15 req/min on free tier
};

// ── Helpers ───────────────────────────────────────────────────────────────
function slug(name) {
  return name.toLowerCase().replace(/[^a-z0-9]+/g, '-').replace(/(^-|-$)/g, '');
}
function loadManifest() {
  try { return JSON.parse(fs.readFileSync(MANIFEST, 'utf8')); }
  catch { return {}; }
}
function saveManifest(m) {
  fs.writeFileSync(MANIFEST, JSON.stringify(m, null, 2));
}
function sleep(ms) { return new Promise(r => setTimeout(r, ms)); }

function buildPrompt(weapon) {
  const cat = weapon.category;
  const stats = weapon.stats || {};
  const cls = stats.class || '';

  const catDesc = {
    'Melee Weapon':  'melee weapon held in gauntleted hands',
    'Ranged Weapon': 'ranged firearm positioned for display',
    'Grenade':       'grenade or explosive device, palm-sized',
    'Armour':        'suit of Space Marine power armor',
    'Other':         'piece of military equipment',
  }[cat] || 'weapon';

  return [
    `Warhammer 40000 Deathwatch ${catDesc}: "${weapon.name}".`,
    cls ? `Type: ${cls}.` : '',
    'Dark gothic sci-fi aesthetic, ultra-detailed digital painting.',
    'Dramatic atmospheric lighting, glowing sigils, battle-worn metal.',
    'Centered isolated subject, black background, no text, no watermark.',
    'Style: official Games Workshop concept art.',
  ].filter(Boolean).join(' ');
}

// ── OpenAI DALL-E 3 ───────────────────────────────────────────────────────
async function generateOpenAI(prompt) {
  const key = process.env.OPENAI_API_KEY;
  if (!key) throw new Error('OPENAI_API_KEY is not set');

  const res = await fetch('https://api.openai.com/v1/images/generations', {
    method: 'POST',
    headers: { 'Content-Type': 'application/json', Authorization: `Bearer ${key}` },
    body: JSON.stringify({ model: 'dall-e-3', prompt, n: 1, size: '1024x1024',
      quality: 'standard', response_format: 'url' }),
  });
  if (!res.ok) throw new Error(`OpenAI error ${res.status}: ${await res.text()}`);
  const data = await res.json();
  const url = data.data[0].url;

  // download PNG
  const imgRes = await fetch(url);
  if (!imgRes.ok) throw new Error(`Download failed: ${imgRes.status}`);
  return imgRes.buffer();
}

// ── Hugging Face FLUX.1-schnell ───────────────────────────────────────────
async function generateHuggingFace(prompt) {
  const key = process.env.HF_TOKEN || process.env.HUGGINGFACE_TOKEN;
  if (!key) throw new Error('HF_TOKEN is not set');

  const res = await fetch(
    'https://router.huggingface.co/hf-inference/models/black-forest-labs/FLUX.1-schnell',
    {
      method: 'POST',
      headers: { 'Content-Type': 'application/json', Authorization: `Bearer ${key}` },
      body: JSON.stringify({ inputs: prompt }),
    }
  );
  if (!res.ok) throw new Error(`HuggingFace error ${res.status}: ${await res.text()}`);
  return res.buffer();
}

// ── Google Gemini Imagen 3 ────────────────────────────────────────────────
async function generateGemini(prompt) {
  const key = process.env.GOOGLE_API_KEY || process.env.GEMINI_API_KEY;
  if (!key) throw new Error('GOOGLE_API_KEY (or GEMINI_API_KEY) is not set');

  const model = 'imagen-3.0-generate-002';
  const url = `https://generativelanguage.googleapis.com/v1beta/models/${model}:predict?key=${key}`;

  const res = await fetch(url, {
    method: 'POST',
    headers: { 'Content-Type': 'application/json' },
    body: JSON.stringify({
      instances: [{ prompt }],
      parameters: {
        sampleCount: 1,
        aspectRatio: '1:1',
        safetyFilterLevel: 'block_only_high',
        personGeneration: 'allow_adult',
      },
    }),
  });

  if (!res.ok) throw new Error(`Gemini error ${res.status}: ${await res.text()}`);
  const data = await res.json();
  const b64 = data.predictions?.[0]?.bytesBase64Encoded;
  if (!b64) throw new Error('No image in Gemini response: ' + JSON.stringify(data));
  return Buffer.from(b64, 'base64');
}

// ── Main ──────────────────────────────────────────────────────────────────
async function main() {
  const generate = PROVIDER === 'openai' ? generateOpenAI
                 : PROVIDER === 'gemini' ? generateGemini
                 : generateHuggingFace;
  const delayMs  = DELAY[PROVIDER] ?? 2_000;

  const keyVar = PROVIDER === 'openai' ? 'OPENAI_API_KEY'
               : PROVIDER === 'gemini' ? 'GOOGLE_API_KEY'
               : 'HF_TOKEN';
  const keyVal = process.env[keyVar] || process.env.GEMINI_API_KEY || process.env.HUGGINGFACE_TOKEN;
  if (!keyVal) {
    console.error(`Error: ${keyVar} environment variable is not set.`);
    process.exit(1);
  }

  console.log(`Provider: ${PROVIDER.toUpperCase()} | Limit: ${LIMIT === Infinity ? 'all' : LIMIT} | Force: ${FORCE}`);

  console.log(`\nFetching weapons from ${WEAPONS_API} …`);
  let weapons;
  try {
    const res = await fetch(WEAPONS_API);
    if (!res.ok) throw new Error(`HTTP ${res.status}`);
    weapons = await res.json();
  } catch (e) {
    console.error('Could not reach server:', e.message);
    console.error('Make sure the backend is running: npm run pm2:start');
    process.exit(1);
  }
  console.log(`Found ${weapons.length} weapons/items`);

  fs.mkdirSync(OUTPUT_DIR, { recursive: true });
  const manifest = loadManifest();

  let generated = 0, skipped = 0, errors = 0;

  for (const weapon of weapons) {
    if (generated >= LIMIT) break;

    const key  = slug(weapon.name);
    const file = `${key}.png`;
    const dest = path.join(OUTPUT_DIR, file);
    const webPath = `/weapon-images/${file}`;

    if (!FORCE && manifest[weapon.id] && fs.existsSync(dest)) {
      skipped++;
      continue;
    }

    const prompt = buildPrompt(weapon);
    console.log(`\n[${generated + 1}] ${weapon.name} (${weapon.category})`);
    console.log(`  prompt: ${prompt.substring(0, 90)}…`);

    try {
      const buf = await generate(prompt);
      fs.writeFileSync(dest, buf);
      manifest[weapon.id] = { name: weapon.name, path: webPath, provider: PROVIDER, generated: new Date().toISOString() };
      saveManifest(manifest);
      console.log(`  ✓ Saved → ${webPath}`);
      generated++;
    } catch (e) {
      console.error(`  ✗ Failed: ${e.message}`);
      errors++;
    }

    // Rate-limit throttle
    const isLast = weapons.indexOf(weapon) >= weapons.length - 1 || generated >= LIMIT;
    if (!isLast) {
      process.stdout.write(`  … waiting ${delayMs / 1000}s …`);
      await sleep(delayMs);
      process.stdout.write('\r\x1b[2K');
    }
  }

  console.log(`\nDone. Generated: ${generated}, Skipped: ${skipped}, Errors: ${errors}`);
  if (generated > 0) console.log(`Images in: public/weapon-images/`);
}

main().catch(e => { console.error(e); process.exit(1); });