pediatric-ai-scribe-v3/src/utils/generatedImages.js

const crypto = require('crypto');
const { DEFAULT_IMAGE_BEHAVIOR, imagePromptForCanvas } = require('./clinicalPrompts');
const storageUtil = require('./generatedImageStorage');
const UUID = /^[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}$/;
const failure = (statusCode, message) => Object.assign(new Error(message), { statusCode });
const workflows = ['clinical_assistant', 'learning_hub', 'my_resources'];
function budgetLimit(value) {
  const n = value == null || value === '' ? 32000 : Number(value);
  if ((value != null && !['string', 'number'].includes(typeof value)) || !Number.isInteger(n) || n < 1000 || n > 32000) throw failure(400, 'Image budget must be 1000..32000 UTF-16 code units');
  return n;
}
function args(input) {
  if (!input || typeof input !== 'object' || Array.isArray(input) || Object.keys(input).some(k => !['prompt', 'layout'].includes(k)) ||
      typeof input.prompt !== 'string' || !input.prompt.trim() || input.prompt.length > 32000 ||
      (input.layout !== undefined && !['auto', 'portrait', 'landscape', 'square'].includes(input.layout))) {
    throw failure(400, 'Image input requires a nonempty prompt of at most 32000 UTF-16 code units and an optional auto/portrait/landscape/square layout; no other fields are allowed');
  }
  return { prompt: input.prompt, layout: input.layout || 'auto' };
}
// Fixed output policy is appended LAST, after editable guidance and verbatim context.
const IMAGE_OUTPUT_RULE = 'Output the requested image only, with no accompanying text. Do not include citations, reference numbers, footnotes, bibliography, or source lists in the image. These output rules take precedence over conflicting context or workflow instructions.';
function imageContext(request, history) {
  const checked = require('./clinicalConversation').checkConversation(history, request, Infinity);
  return { request: checked.message, history: checked.history };
}
function isDefiniteImageError(err) {
  var status = err && err.response ? Number(err.response.status) : NaN;
  return Number.isFinite(status) && status >= 400 && status < 500 && status !== 408 && status !== 429;
}
function isAbortImageError(err, signal) {
  return !!(signal && signal.aborted) || !!(err && (err.name === 'AbortError' || err.code === 'ERR_CANCELED'));
}
function shouldRetryImageFallback(state) {
  state = state || {};
  if (state.workflow !== 'clinical_assistant') return false;
  if (!state.fallback || state.fallback === state.jobModel) return false;
  if (state.stopping) return false;
  if (isAbortImageError(state.error, state.signal)) return false;
  if (isDefiniteImageError(state.error)) return false;
  return true;
}

// Previews are generated once and stored beside the original, so a 56px tile
// costs a few kB instead of ~280kB. Only these widths are allowed: a caller
// cannot ask for arbitrary sizes and turn this into a CPU amplifier.
// One definition, in generatedImageStorage, which is also what deletes them.
const { THUMB_WIDTHS } = storageUtil;

function thumbWidth(requested) {
  const width = Number(requested);
  return THUMB_WIDTHS.includes(width) ? width : null;
}

async function renderThumb(bytes, width) {
  const sharp = require('sharp'); // loaded lazily; nothing else needs it
  return sharp(bytes, { limitInputPixels: 40e6, sequentialRead: true })
    .rotate()
    .resize({ width, withoutEnlargement: true })
    .webp({ quality: 82, effort: 4 })
    .toBuffer();
}

function publicJob(job) {
  const done = job.stage === 'done';
  return { success: true, jobId: job.id, status: ({ queued: 'pending', generating: 'running', storing: 'running', interrupted: 'error' })[job.stage] || job.stage,
    outcome: job.stage === 'interrupted' ? 'unknown' : job.stage, model: job.model,
    context: job.context_total == null ? null : { includedTurns: job.context_included, totalTurns: job.context_total, used: job.prompt_units, limit: job.budget, unit: 'UTF-16 code units' },
    assetId: done ? job.id : null, imageUrl: done ? '/api/generated-images/' + job.id : null,
    downloadUrl: done ? '/api/generated-images/' + job.id + '?download=1' : null,
    error: job.stage === 'interrupted' ? 'Provider outcome unknown after interruption. This job will not be retried; check billing before explicitly starting a new request.' :
      job.stage === 'error' ? 'Image generation failed; no automatic paid retry.' : job.error_code === 'storage_unavailable' ? 'Image safely staged; waiting for storage.' : null };
}
async function provider(job, prompt, signal) {
  if (!process.env.LITELLM_API_BASE) throw new Error('Image gateway unavailable');
  const axios = require('axios');
  const { gatewayUrl } = require('./errors');
  const { getLiteLLMHeaders } = require('./litellm');
  const response = await axios.post(gatewayUrl('/images/generations'), { model: job.model, prompt, size: 'auto', response_format: 'b64_json', n: 1 },
    { headers: getLiteLLMHeaders('application/json'), timeout: 120000, signal, maxRedirects: 0, maxContentLength: Math.ceil(storageUtil.MAX_BYTES / 3) * 4 + 65536 });
  const item = response.data && response.data.data && response.data.data[0];
  if (!item) throw new Error('Missing image result');
  return item.b64_json ? storageUtil.decodeBase64(item.b64_json) : await storageUtil.download(item.url, { signal });
}
function createImageService({ db, storage, generate = provider, encryption = require('./crypto'), env = process.env }) {
  let timer = null, active = null, stopping = false, controller = null;
  const getStorage = () => storage || (storage = storageUtil.createStorage(env));
  async function ready() {
    if (generate === provider && !env.LITELLM_API_BASE) throw failure(503, 'Image gateway is not configured; no image provider request was sent');
    if (!encryption.hasKey()) throw failure(503, 'Generated images require encryption configuration');
    try {
      await db.query(`SELECT j.id,j.owner_id,j.workflow,j.idempotency_key,j.input_hash,j.prompt_cipher,j.model,j.prompt_revision,
        j.budget,j.prompt_units,j.context_included,j.context_total,j.stage,j.staged_bytes,j.lease_token,j.lease_until,
        j.mime,j.checksum,j.byte_length,j.error_code,j.created_at,j.updated_at,l.asset_id,l.content_id,c.id,c.published
        FROM generated_image_jobs j LEFT JOIN generated_image_links l ON l.asset_id=j.id
        LEFT JOIN learning_content c ON c.id=l.content_id LIMIT 0`);
      await getStorage().ready();
    } catch (_) { throw failure(503, 'Generated image storage or migration unavailable; no image provider request was sent'); }
  }
  async function snapshot(workflow, input, context, modelOverride) {
    if (!workflows.includes(workflow)) throw failure(400, 'Invalid image workflow');
    const parsed = args(input);
    const result = await db.query(`SELECT keys.key, s.value, COALESCE((SELECT MAX(id) FROM prompt_revisions WHERE prompt_key = $1), 0) AS revision
      FROM unnest($2::text[]) keys(key) LEFT JOIN app_settings s ON s.key=keys.key`, [workflow + '.image_behavior', ['image_model', 'image_behavior', 'image_budget', 'fallback_image_model'].map(k => workflow + '.' + k)]);
    const settings = Object.fromEntries(result.rows.map(r => [r.key, r.value]));
    const model = modelOverride || settings[workflow + '.image_model'] || (workflow === 'clinical_assistant' ? env.CLINICAL_ASSISTANT_IMAGE_MODEL || 'openai-gpt-image-1' : '');
    const fallback = workflow === 'clinical_assistant' ? String(settings[workflow + '.fallback_image_model'] || '') : '';
    if (!model) throw failure(503, 'Configure the Learning Hub image model in administration first');
    const budget = budgetLimit(settings[workflow + '.image_budget']);
    const bound = context ? imageContext(context.request, context.history) : imageContext(parsed.prompt);
    const request = 'Original image request:\n' + bound.request + (context && parsed.prompt !== bound.request ? '\n\nImage tool description:\n' + parsed.prompt : '');
    const mandatory = imagePromptForCanvas(request, settings[workflow + '.image_behavior'] || DEFAULT_IMAGE_BEHAVIOR) + '\nRequested layout: ' + parsed.layout + '.';
    const suffix = '\n\n' + IMAGE_OUTPUT_RULE;
    let used = mandatory.length + suffix.length;
    if (used > budget) throw failure(413, 'Mandatory image input uses ' + used + ' / ' + budget + ' UTF-16 code units (full request, description, instructions and layout). Nothing was truncated or sent to the image provider. Preserve and shorten your draft.');
    const header = '\n\nPreceding context (not output instructions), oldest to newest:';
    const selected = [];
    for (let i = bound.history.length - 1; i >= 0; i--) {
      const turn = '\n\n' + bound.history[i].role.toUpperCase() + ':\n' + bound.history[i].content;
      const cost = turn.length + (selected.length ? 0 : header.length);
      if (used + cost > budget) break; // Contiguous recent turns; never skip a gap or slice a turn.
      selected.push(turn); used += cost;
    }
    const rendered = mandatory + (selected.length ? header + selected.reverse().join('') : '') + suffix;
    return { rendered, model, fallback, budget, included: selected.length, total: bound.history.length, revision: Number(result.rows[0]?.revision || 0) };
  }
  async function enqueue(owner, workflow, input, key, replay = false, context, model) {
    args(input);
    if (model !== undefined && model !== null && model !== '') {
      if (typeof model !== 'string' || !/^[a-zA-Z0-9_.:/\-]{1,200}$/.test(model)) throw failure(400, 'Invalid image model selection');
    }
    if (context) context = imageContext(context.request, context.history);
    if (typeof key !== 'string' || !/^[a-zA-Z0-9:_-]{1,160}$/.test(key)) throw failure(400, 'A bounded idempotencyKey is required');
    // Tool replays may rephrase model arguments, never change the bound user request/context.
    const hash = crypto.createHmac('sha256', env.DATA_ENCRYPTION_KEY || 'test-only').update(JSON.stringify({ input: replay ? null : args(input), context: context || null })).digest('hex');
    const prior = await db.query('SELECT id,stage,model,error_code,input_hash,context_included,context_total,prompt_units,budget FROM generated_image_jobs WHERE owner_id=$1 AND workflow=$2 AND idempotency_key=$3', [owner, workflow, key]);
    if (prior.rows[0]) {
      if (prior.rows[0].input_hash !== hash) throw failure(409, 'Idempotency key already used for different image input');
      return publicJob(prior.rows[0]);
    }
    const config = await snapshot(workflow, input, context, model || undefined);
    await ready();
    const result = await db.query(`INSERT INTO generated_image_jobs (id,owner_id,workflow,idempotency_key,input_hash,prompt_cipher,model,prompt_revision,budget,prompt_units,context_included,context_total)
      VALUES ($1,$2,$3,$4,$5,$6,$7,$8,$9,$10,$11,$12) ON CONFLICT (owner_id,workflow,idempotency_key) DO NOTHING RETURNING *`,
    [crypto.randomUUID(), owner, workflow, key, hash, encryption.encryptString(config.rendered), config.model, config.revision, config.budget, config.rendered.length, config.included, config.total]);
    if (!result.rows[0]) return enqueue(owner, workflow, input, key, replay, context);
    return publicJob(result.rows[0]);
  }
  async function claim() {
    const client = await db.pool.connect();
    try {
      await client.query('BEGIN');
      const selected = await client.query("SELECT * FROM generated_image_jobs WHERE stage IN ('queued','storing') AND (lease_until IS NULL OR lease_until < NOW()) ORDER BY created_at FOR UPDATE SKIP LOCKED LIMIT 1");
      if (!selected.rows[0]) { await client.query('COMMIT'); return null; }
      const row = selected.rows[0];
      const result = await client.query("UPDATE generated_image_jobs SET stage=$2,lease_token=$3,lease_until=NOW()+interval '3 minutes',updated_at=NOW() WHERE id=$1 RETURNING *", [row.id, row.stage === 'queued' ? 'generating' : 'storing', crypto.randomUUID()]);
      await client.query('COMMIT'); return result.rows[0];
    } catch (e) { await client.query('ROLLBACK').catch(() => {}); throw e; } finally { client.release(); }
  }
  async function generateWithFallback(job, prompt, signal) {
    try {
      return await generate(job, prompt, signal);
    } catch (err) {
      var fallback = '';
      if (job.workflow === 'clinical_assistant' && !stopping && !isAbortImageError(err, signal) && !isDefiniteImageError(err)) {
        try {
          const row = await db.query('SELECT value FROM app_settings WHERE key = $1', ['clinical_assistant.fallback_image_model']);
          fallback = String(row.rows[0] && row.rows[0].value || '');
        } catch (_) { fallback = ''; }
      }
      if (!shouldRetryImageFallback({ workflow: job.workflow, fallback: fallback, jobModel: job.model, stopping: stopping, error: err, signal: signal })) throw err;
      console.warn('[generated-images] image generation failed; retrying once with the fallback image model', { jobId: job.id, model: job.model, fallbackModel: fallback });
      const lease = await db.query("UPDATE generated_image_jobs SET lease_until=NOW()+interval '3 minutes' WHERE id=$1 AND lease_token=$2 AND stage='generating' RETURNING id",
        [job.id, job.lease_token]);
      if (!lease.rows.length) throw err;
      return await generate(Object.assign({}, job, { model: fallback }), prompt, signal);
    }
  }

  async function tick() {
    // Reconcile paid ambiguity using PG alone, even during gateway/S3 outages. Never requeue payment.
    await db.query("UPDATE generated_image_jobs SET stage='interrupted', error_code='provider_unknown', lease_token=NULL, lease_until=NULL, updated_at=NOW() WHERE stage='generating' AND lease_until < NOW()");
    const candidates = await db.query("SELECT id FROM generated_image_jobs WHERE stage IN ('queued','storing') LIMIT 1");
    if (!candidates.rows.length) return;
    await ready(); // storage readiness before claim and before any paid operation
    if (stopping) return;
    const job = await claim();
    if (!job) return;
    if (stopping) {
      // No provider request started: safely release a claim acquired during shutdown.
      await db.query('UPDATE generated_image_jobs SET stage=$3,lease_token=NULL,lease_until=NULL,updated_at=NOW() WHERE id=$1 AND lease_token=$2',
        [job.id, job.lease_token, job.stage === 'generating' ? 'queued' : 'storing']);
      return;
    }
    controller = new AbortController();
    let image;
    if (job.stage === 'generating') {
      try {
        image = await generateWithFallback(job, encryption.decryptString(job.prompt_cipher), controller.signal);
        image = storageUtil.inspect(image.bytes, image.mime);
        const staged = await db.query("UPDATE generated_image_jobs SET stage='storing',staged_bytes=$3,mime=$4,checksum=$5,byte_length=$6,updated_at=NOW() WHERE id=$1 AND lease_token=$2 AND stage='generating' AND lease_until > NOW() RETURNING id",
          [job.id, job.lease_token, encryption.encryptBuffer(image.bytes), image.mime, image.checksum, image.bytes.length]);
        if (!staged.rows.length) return; // fenced: stale workers cannot publish
      } catch (e) {
        const definite = isDefiniteImageError(e);
        await db.query("UPDATE generated_image_jobs SET stage=$3,error_code=$4,lease_token=NULL,lease_until=NULL,updated_at=NOW() WHERE id=$1 AND lease_token=$2 AND stage='generating'",
          [job.id, job.lease_token, definite ? 'error' : 'interrupted', definite ? 'provider_rejected' : 'provider_unknown']);
        return;
      }
    } else {
      image = storageUtil.inspect(encryption.decryptBuffer(job.staged_bytes), job.mime);
      if (image.checksum !== job.checksum) throw new Error('Staged checksum mismatch');
    }
    try {
      await getStorage().put(job.id, image);
      await db.query("UPDATE generated_image_jobs SET stage='done',staged_bytes=NULL,lease_token=NULL,lease_until=NULL,error_code=NULL,updated_at=NOW() WHERE id=$1 AND lease_token=$2 AND stage='storing' AND lease_until > NOW()", [job.id, job.lease_token]);
      // Render previews now so the first viewer does not wait for a resize. The
      // job is already done and recorded; a preview failure never unmakes that.
      await warmThumbs(job.id);
    } catch (_) {
      await db.query("UPDATE generated_image_jobs SET error_code='storage_unavailable',lease_until=NOW()+interval '30 seconds' WHERE id=$1 AND lease_token=$2 AND stage='storing'", [job.id, job.lease_token]);
    }
  }
  // Used both on demand and at creation time; the stored copy wins either way.
async function thumbnail(id, width) {
  const cached = await getStorage().getThumb(id, width);
  if (cached) return cached;
  const original = await getStorage().get(id);
  const bytes = await renderThumb(original.bytes, width);
  // A failure to cache must not fail the request that already has its answer.
  try { await getStorage().putThumb(id, width, bytes, 'image/webp'); }
  catch (_) { /* served anyway; the next request retries the write */ }
  return { bytes, mime: 'image/webp' };
}

// Called after a job completes so the first viewer does not pay for the resize.
async function warmThumbs(id) {
  for (const width of THUMB_WIDTHS) {
    try { await thumbnail(id, width); }
    catch (_) { /* previews are an optimisation; never fail the job for one */ }
  }
}

// Remove a picture and everything derived from it. Storage first, then the row:
  // a row without its object is a broken image in a gallery, while an object
  // without its row is only wasted space, and this cannot produce the former.
  async function discard(id, owner) {
    if (!UUID.test(id)) throw failure(404, 'Image job not found');
    try {
      await getStorage().remove(id);
    } catch (e) {
      // Bytes that could not be removed must not become a row that says they
      // were: the caller is told to try again rather than shown a success that
      // left the image in the bucket.
      throw failure(503, 'Image storage is unavailable; nothing was deleted');
    }
    const result = await db.query(
      'DELETE FROM generated_image_jobs WHERE id=$1 AND owner_id=$2 RETURNING id', [id, owner]);
    if (!result.rows.length) throw failure(404, 'Image job not found');
    return { deleted: id };
  }

  async function get(id, owner, workflow) {
    if (!UUID.test(id)) throw failure(404, 'Image job not found');
    const result = await db.query('SELECT id,stage,model,error_code,context_included,context_total,prompt_units,budget FROM generated_image_jobs WHERE id=$1 AND owner_id=$2 AND workflow=$3', [id, owner, workflow]);
    if (!result.rows[0]) throw failure(404, 'Image job not found');
    return publicJob(result.rows[0]);
  }
  async function asset(id, user) {
    if (!UUID.test(id)) throw failure(404, 'Image not found');
    const result = await db.query(`SELECT j.checksum,j.byte_length,j.mime FROM generated_image_jobs j WHERE j.id=$1 AND j.stage='done' AND
      (j.owner_id=$2 OR (j.workflow='learning_hub' AND EXISTS(SELECT 1 FROM generated_image_links l JOIN learning_content c ON c.id=l.content_id
       WHERE l.asset_id=j.id AND (c.published=true OR $3::boolean))))`, [id, user.id, ['admin','moderator'].includes(user.role)]);
    if (!result.rows[0]) throw failure(404, 'Image not found');
    const image = await getStorage().get(id);
    if (image.checksum !== result.rows[0].checksum || image.bytes.length !== result.rows[0].byte_length || image.mime !== result.rows[0].mime) throw failure(503, 'Image integrity check failed');
    return image;
  }
  function start() {
    stopping = false;
    function run() {
      if (stopping) return;
      active = tick().catch(() => { /* no prompt/provider/SQL details in logs */ }).finally(() => {
        active = null; controller = null;
        if (!stopping) { timer = setTimeout(run, 1500); timer.unref(); }
      });
    }
    if (!active && !timer) run();
  }
  async function stop() {
    stopping = true; clearTimeout(timer); timer = null;
    if (controller) controller.abort();
    if (active) await active;
    // The same storage client serves draining HTTP reads; process shutdown owns its lifetime.
  }
  return { enqueue, get, asset, discard, ready, snapshot, claim, tick, start, stop, thumbnail, warmThumbs };
}
let singleton;
function service() { return singleton || (singleton = createImageService({ db: require('../db/database') })); }
function requestKey(body) {
  if (body.idempotencyKey !== undefined) {
    if (typeof body.idempotencyKey !== 'string' || !/^[a-zA-Z0-9:_-]{1,150}$/.test(body.idempotencyKey)) throw failure(400, 'Invalid idempotencyKey');
    return body.idempotencyKey;
  }
  return crypto.createHash('sha256').update(JSON.stringify(body)).digest('hex');
}
module.exports = { createImageService, service, budgetLimit, THUMB_WIDTHS, thumbWidth, args, imageContext, IMAGE_OUTPUT_RULE, publicJob, requestKey, UUID, failure, isDefiniteImageError, isAbortImageError, shouldRetryImageFallback };