atelier/web/Generate/useGenerateManager.ts

import { useEffect, useState } from 'react';
import { useAtelierRefresh } from 'officerdev';
import { useAtelierClient } from '../lib/api';
import type { JobResponse, ProviderStatus } from '../lib/api';

export const SIZES = ['auto', '1024x1024', '1536x1024', '1024x1536'] as const;
export const QUALITIES = ['auto', 'low', 'medium', 'high'] as const;
export const BACKGROUNDS = ['auto', 'transparent', 'opaque'] as const;

/**
 * The form. Providers are passed IN rather than fetched here — `useProviders` owns them, because the
 * dropdown above this form and the model list inside it have to be the same list. They were two
 * fetches, and adding a provider left the form showing the old models until the page was reloaded.
 */
export const useGenerateManager = (providers: ProviderStatus[]) => {
  const client = useAtelierClient();
  const [, bumpRefresh] = useAtelierRefresh();

  const [model, setModel] = useState('');
  const [prompt, setPrompt] = useState('');
  const [n, setN] = useState(1);
  const [size, setSize] = useState<string>('auto');
  const [quality, setQuality] = useState<string>('auto');
  const [background, setBackground] = useState<string>('auto');

  const [job, setJob] = useState<JobResponse | null>(null);
  const [running, setRunning] = useState(false);
  const [error, setError] = useState<string | null>(null);

  // Follow the provider list: pick the default's model, and only overwrite a choice the user made if
  // the provider behind it has gone away.
  useEffect(() => {
    const usable = providers.filter((row) => row.configured && row.enabled);
    if (usable.length === 0) return;
    if (model && usable.some((row) => row.model === model)) return;
    setModel((usable.find((row) => row.is_default) ?? usable[0]!).model);
    // eslint-disable-next-line react-hooks/exhaustive-deps
  }, [providers]);

  // Derived in the hook, not in the component.
  const usableProviders = providers.filter((provider) => provider.configured && provider.enabled);
  // One model per configured provider, deduped: two providers on the same model is legitimate (two
  // keys, two accounts) and the request picks between them by model alone.
  const models = [...new Set(usableProviders.map((provider) => provider.model))];
  const ready = Boolean(prompt.trim()) && Boolean(model) && !running;

  async function generate() {
    if (!ready) return;

    setRunning(true);
    setError(null);
    setJob(null);

    try {
      // Only what the user actually chose. `auto` is the provider's own default, so sending it would
      // be inventing a parameter they did not set.
      const body: Record<string, unknown> = { prompt: prompt.trim(), model, n };
      if (size !== 'auto') body.size = size;
      if (quality !== 'auto') body.quality = quality;
      if (background !== 'auto') body.background = background;

      const { id } = await client.startGeneration(body);

      // Long-poll in a loop: each call waits up to 30s server-side, so this is one request per half
      // minute rather than a busy poll, and it survives a generation that takes longer.
      for (;;) {
        const result = await client.awaitGeneration(id, 30);
        if (result.atelier && ['succeeded', 'failed', 'partial'].includes(result.atelier.status)) {
          setJob(result);
          bumpRefresh();
          return;
        }
        setJob(result);
      }
    } catch (err) {
      setError(err instanceof Error ? err.message : String(err));
    } finally {
      setRunning(false);
    }
  }

  return {
    providers,
    usableProviders,
    models,
    model,
    setModel,
    prompt,
    setPrompt,
    n,
    setN,
    size,
    setSize,
    quality,
    setQuality,
    background,
    setBackground,
    job,
    running,
    error,
    ready,
    generate,
  };
};

export type GenerateManager = ReturnType<typeof useGenerateManager>;