tests/test_connections/test_endpoints.mjs

/* eslint-disable no-console */
import 'dotenv/config';

const STORY_BASE = (process.env.AI_STORY_API_GRADIO_URL || '').replace(/\/+$/, '');
const STORY_SECRET = process.env.AI_STORY_API_SECRET_TOKEN || '';

const IMG_BASE = (process.env.AI_FAST_IMAGE_SERVER_API_GRADIO_URL || '').replace(/\/+$/, '');
const IMG_SECRET = process.env.AI_FAST_IMAGE_SERVER_API_SECRET_TOKEN || '';

const HF_TOKEN = process.env.HF_TOKEN || process.env.HUGGING_FACE_HUB_TOKEN || '';

const FN_INDEX = Number(process.env.AI_STORY_API_FN_INDEX ?? 0); // legacy only
const DEBUG = (process.env.DEBUG_STORY_API || '').toLowerCase() === 'true';

function abbrev(s, n = 1800) {
  if (s == null) return String(s);
  const str = typeof s === 'string' ? s : JSON.stringify(s);
  return str.length > n ? str.slice(0, n) + '…' : str;
}

function assertEnv() {
  if (!STORY_BASE) throw new Error('Missing AI_STORY_API_GRADIO_URL in .env');
  if (!STORY_SECRET) throw new Error('Missing AI_STORY_API_SECRET_TOKEN in .env');
}

function authHeaders() {
  const h = {};
  if (HF_TOKEN) h.Authorization = `Bearer ${HF_TOKEN}`;
  return h;
}

async function withTimeout(promise, ms = 120_000) {
  return Promise.race([
    promise,
    new Promise((_, rej) => setTimeout(() => rej(new Error(`Timeout after ${ms}ms`)), ms)),
  ]);
}

async function timedFetch(url, init) {
  const t0 = Date.now();
  const res = await withTimeout(fetch(url, init));
  const ms = Date.now() - t0;
  return { res, ms };
}

function stripToJson(text) {
  if (!text) return '';
  const ixBrace = text.indexOf('{');
  const ixBracket = text.indexOf('[');
  const pick = [ixBrace, ixBracket].filter((i) => i >= 0);
  if (!pick.length) return text;
  const i = Math.min(...pick);
  return text.slice(i);
}

/** Parse SSE and return the JSON payload of the **complete** event (throw on error). */
async function readSSEComplete(res, { maxMs = 120_000 } = {}) {
  const ct = (res.headers.get('content-type') || '').toLowerCase();
  // Some deployments reply JSON directly (no SSE) — accept that.
  if (ct.includes('application/json')) {
    const txt = await res.text();
    try {
      return JSON.parse(stripToJson(txt));
    } catch {
      return null;
    }
  }

  // Proper SSE parse: look for event: complete / event: error
  const reader = res.body.getReader();
  const decoder = new TextDecoder();
  let buf = '';
  let currEvent = null;
  let dataLines = [];
  const started = Date.now();

  const handleBlock = () => {
    if (!currEvent) return;
    const payloadText = dataLines.join('\n');
    if (DEBUG) console.log(`   [SSE event=${currEvent}] ${abbrev(payloadText, 600)}`);
    if (currEvent === 'complete') {
      if (!payloadText) return { type: 'complete', json: null };
      try {
        return { type: 'complete', json: JSON.parse(stripToJson(payloadText)) };
      } catch {
        return { type: 'complete', json: null };
      }
    }
    if (currEvent === 'error') {
      let detail = null;
      try {
        detail = payloadText ? JSON.parse(stripToJson(payloadText)) : null;
      } catch {
        detail = payloadText;
      }
      return { type: 'error', detail };
    }
    return undefined;
  };

  while (true) {
    if (Date.now() - started > maxMs) throw new Error(`SSE read timeout after ${maxMs}ms`);
    const { value, done } = await reader.read();
    if (done) break;
    buf += decoder.decode(value, { stream: true });

    // Normalize CRLF and split into blocks
    let idx;
    while ((idx = buf.search(/\r?\n\r?\n/)) !== -1) {
      const raw = buf.slice(0, idx);
      buf = buf.slice(idx).replace(/^\r?\n/, '');
      const lines = raw.split(/\r?\n/);

      // Parse fields inside this SSE block
      currEvent = null;
      dataLines = [];
      for (const line of lines) {
        if (!line) continue;
        // comment keep-alive lines start with ':'
        if (line[0] === ':') continue;
        const c = line.indexOf(':');
        const field = (c === -1 ? line : line.slice(0, c)).trim();
        const val = (c === -1 ? '' : line.slice(c + 1)).replace(/^\s/, '');
        if (field === 'event') currEvent = val;
        else if (field === 'data') dataLines.push(val);
      }

      const outcome = handleBlock();
      if (outcome?.type === 'complete') return outcome.json;
      if (outcome?.type === 'error') {
        const msg = typeof outcome.detail === 'string'
          ? outcome.detail
          : JSON.stringify(outcome.detail);
        throw new Error(`SSE error event: ${abbrev(msg ?? 'null', 1200)}`);
      }
    }
  }
  return null;
}

/** Gradio v5: POST /gradio_api/call/<fn> -> {event_id}, then GET /gradio_api/call/<fn>/<event_id> (SSE) */
async function tryCallApi(base, fnName, dataArray) {
  const callUrl = `${base}/gradio_api/call/${fnName}`;
  const postBody = { data: dataArray };

  console.log(`\nPOST ${callUrl}`);
  if (DEBUG) console.log('  body:', abbrev(JSON.stringify(postBody), 600));

  const { res: postRes, ms: postMs } = await timedFetch(callUrl, {
    method: 'POST',
    headers: { 'Content-Type': 'application/json', Accept: 'application/json', ...authHeaders() },
    body: JSON.stringify(postBody),
    cache: 'no-store',
  });

  const postText = await postRes.text();
  const postJson = (() => {
    try {
      return JSON.parse(stripToJson(postText));
    } catch {
      return null;
    }
  })();

  if (!postRes.ok) {
    const detail = postJson?.detail ?? postJson?.error ?? postJson?.message ?? postText ?? '(empty body)';
    throw new Error(`HTTP ${postRes.status} ${postRes.statusText} in ${postMs}ms — ${abbrev(String(detail), 1200)}`);
  }

  const eventId = postJson?.event_id || postJson?.eventId || postJson?.eventID;
  if (!eventId) {
    console.error('  ⚠️ POST response had no event_id. Raw body:\n', abbrev(postText, 1200));
    throw new Error('Missing event_id in call API response');
  }

  const getUrl = `${base}/gradio_api/call/${fnName}/${eventId}`;
  console.log(`GET  ${getUrl}`);

  const { res: getRes, ms: getMs } = await timedFetch(getUrl, {
    method: 'GET',
    headers: { Accept: 'text/event-stream', ...authHeaders() },
    cache: 'no-store',
  });

  if (getRes.status !== 200) {
    const txt = await getRes.text();
    throw new Error(`HTTP ${getRes.status} ${getRes.statusText} in ${getMs}ms — ${abbrev(txt, 800)}`);
  }

  const json = await readSSEComplete(getRes);
  console.log(`   ✅ OK 200 in ${getMs}ms`);
  if (DEBUG) console.log('  json:', abbrev(JSON.stringify(json), 2000));
  return { ok: true, url: getUrl, json, ms: getMs };
}

/** Legacy JSON endpoints (Gradio 3.x/4 compat) */
async function tryPredictLegacy(base, body) {
  const endpoints = [`${base}/api/predict`, `${base}/run/predict`];

  for (const url of endpoints) {
    try {
      console.log(`\nPOST ${url}`);
      if (DEBUG) console.log('  body:', abbrev(JSON.stringify(body), 600));

      const { res, ms } = await timedFetch(url, {
        method: 'POST',
        headers: { 'Content-Type': 'application/json', Accept: 'application/json', ...authHeaders() },
        body: JSON.stringify(body),
        cache: 'no-store',
      });

      const text = await res.text();
      let json = null;
      try {
        json = JSON.parse(stripToJson(text));
      } catch {}
      if (res.ok) {
        console.log(`   ✅ OK ${res.status} in ${ms}ms`);
        if (DEBUG) console.log('  json:', abbrev(JSON.stringify(json), 2000));
        return { ok: true, url, json, ms };
      }

      const detail = json?.detail ?? json?.error ?? json?.message ?? text ?? '(empty body)';
      console.error(`   ❌ HTTP ${res.status} ${res.statusText} in ${ms}ms — ${abbrev(String(detail), 1200)}`);
    } catch (e) {
      console.error(`   ❌ ${url} failed: ${e?.message || e}`);
    }
  }
  return { ok: false };
}

/** Normalize Gradio payload into an array of lines [{text, audio}, ...] */
function extractStoryLines(payload) {
  if (!payload) return null;
  let d = payload.data ?? payload;
  if (Array.isArray(d) && Array.isArray(d[0])) d = d[0];
  return Array.isArray(d) ? d : null;
}

async function testStoryEndpoint() {
  assertEnv();
  console.log('\n== Story endpoint check ==');
  console.log('BASE :', STORY_BASE);
  console.log('TOKEN:', STORY_SECRET ? '(present)' : '(MISSING)');
  console.log('FN   :', FN_INDEX);

  // HEAD liveness (with auth if provided)
  try {
    const { res, ms } = await timedFetch(STORY_BASE + '/', { method: 'HEAD', headers: authHeaders() });
    console.log(`HEAD ${STORY_BASE}/ -> ${res.status} in ${ms}ms`);
  } catch (e) {
    console.warn('HEAD failed:', e?.message || e);
  }

  const voice = 'Cloée';
  const prompt = 'A short wholesome bedtime story about a friendly dragon and a village.';

  // 1) Gradio v5 call API (SSE)
  try {
    const out = await tryCallApi(STORY_BASE, 'predict', [STORY_SECRET, prompt, voice]);
    console.log('Story call API response:', abbrev(JSON.stringify(out.json), 2000));

    const lines = extractStoryLines(out.json);
    if (Array.isArray(lines)) {
      console.log(`   ✅ Received ${lines.length} story lines`);
      return;
    }

    console.warn('⚠️ Story endpoint returned no recognized payload (SSE complete missing/empty). Treating as reachable.');
    return;
  } catch (e) {
    console.warn('Gradio v5 call API failed for story:', e?.message || e);
  }

  // 2) Legacy fallback
  const legacy = await tryPredictLegacy(STORY_BASE, { fn_index: FN_INDEX, data: [STORY_SECRET, prompt, voice] });
  if (!legacy.ok) throw new Error('Story endpoint failed (both call API and legacy).');
  const lines = extractStoryLines(legacy.json);
  if (!Array.isArray(lines)) {
    console.warn('⚠️ Legacy story payload not recognized, but endpoint responded. Treating as reachable.');
  } else {
    console.log(`   ✅ Legacy story returned ${lines.length} lines`);
  }
}

async function testImageEndpoint() {
  if (!IMG_BASE || !IMG_SECRET) {
    console.log('\n== Image endpoint check skipped (missing IMG envs) ==');
    return;
  }
  console.log('\n== Image endpoint check ==');
  console.log('BASE :', IMG_BASE);
  console.log('TOKEN:', IMG_SECRET ? '(present)' : '(MISSING)');

  // Matches server order: [prompt, negative, seed, width, height, guidance, steps, secret_token]
  const imgData = ['a cute cat sticker', '', 1, 512, 512, 0.0, 4, IMG_SECRET];

  // 1) Gradio v5 call API (SSE)
  try {
    const out = await tryCallApi(IMG_BASE, 'generate', imgData);
    console.log('Image call API response:', abbrev(JSON.stringify(out.json), 2000));

    const payload = out.json?.data ?? out.json;
    const ok =
      Array.isArray(payload) ||
      (payload && typeof payload === 'object' && ('url' in payload || 'path' in payload));

    if (ok) {
      console.log('   ✅ Image endpoint responded.');
      return;
    }
    console.warn('⚠️ Image payload not recognized, but endpoint responded. Treating as reachable.');
    return;
  } catch (e) {
    console.warn('Gradio v5 call API failed for image:', e?.message || e);
  }

  // 2) Legacy fallback
  const legacy = await tryPredictLegacy(IMG_BASE, { fn_index: 0, data: imgData });
  if (!legacy.ok) throw new Error('Image endpoint failed (both call API and legacy).');
  const payload = legacy.json?.data ?? legacy.json;
  const ok =
    Array.isArray(payload) ||
    (payload && typeof payload === 'object' && ('url' in payload || 'path' in payload));
  if (!ok) console.warn('⚠️ Legacy image payload not recognized, but endpoint responded. Treating as reachable.');
  else console.log('   ✅ Legacy image endpoint responded.');
}

(async () => {
  try {
    await testStoryEndpoint();
    await testImageEndpoint();
    console.log('\nAll tests done ✅');
    process.exit(0);
  } catch (e) {
    console.error('\nTests failed ❌:', e?.message || e);
    process.exit(1);
  }
})();